add source code for 支付系统

Author: frankli0324

Challenges/Web/支付系统/server.py

@@ -0,0 +1,130 @@
+import os
+import uuid
+from quart import Quart, render_template, redirect, jsonify, request, session
+from hashlib import pbkdf2_hmac
+from enum import IntEnum
+from tortoise import fields
+from tortoise.models import Model
+from tortoise.contrib.quart import register_tortoise
+from httpx import AsyncClient
+
+app = Quart(__name__)
+app.secret_key = os.urandom(16)
+
+
+class TransactionStatus(IntEnum):
+    SUCCESS = 0
+    PENDING = 1
+    FAILED = 2
+    TIMEOUT = 3
+
+
+class Transaction(Model):
+    id = fields.IntField(pk=True)
+    user = fields.UUIDField()
+    amount = fields.IntField()
+    status = fields.IntEnumField(TransactionStatus)
+    desc = fields.TextField()
+    hash = fields.CharField(64, null=True)
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        for k, v in kwargs.items():
+            self.__setattr__(k, v)
+
+
+async def do_callback(transaction: Transaction):
+    async with AsyncClient() as ses:
+        transaction.status = int(TransactionStatus.FAILED)
+        data = (
+            f'{transaction.id}'
+            f'{transaction.user}'
+            f'{transaction.amount}'
+            f'{transaction.status}'
+            f'{transaction.desc}'
+        ).encode()
+        await ses.post(f'http://localhost:8000/callback', data={
+            'id': transaction.id,
+            'user': transaction.user,
+            'amount': transaction.amount,
+            'desc': transaction.desc,
+            'status': transaction.status,
+            'hash': pbkdf2_hmac('sha256', data, app.secret_key, 2**20).hex()
+        })
+
+
+@app.before_request
+async def create_session():
+    if 'uid' not in session:
+        session['uid'] = str(uuid.uuid4())
+    session['balance'] = 0
+    for tr in await Transaction.filter(user=session['uid']).all():
+        if tr.status == TransactionStatus.SUCCESS:
+            session['balance'] += tr.amount
+
+
+@app.route('/pay')
+async def pay():
+    transaction = await Transaction.create(
+        amount=request.args.get('amount'),
+        desc=request.args.get('desc'),
+        status=TransactionStatus.PENDING,
+        user=uuid.UUID(session.get('uid'))
+    )
+    app.add_background_task(do_callback, transaction)
+    return redirect(f'/transaction?id={transaction.id}')
+
+
+@app.route('/callback', methods=['POST'])
+async def callback():
+    form = dict(await request.form)
+    data = (
+        f'{form.get("id")}'
+        f'{form.get("user")}'
+        f'{form.get("amount")}'
+        f'{form.get("status")}'
+        f'{form.get("desc")}'
+    ).encode()
+    k = pbkdf2_hmac('sha256', data, app.secret_key, 2**20).hex()
+    tr = await Transaction.get(id=int(form.pop('id')))
+    if k != form.get("hash"):
+        return '403'
+    form['status'] = TransactionStatus(int(form.pop('status')))
+    tr.update_from_dict(form)
+    await tr.save()
+    return 'ok'
+
+
+@app.route('/transaction')
+async def transaction():
+    if 'id' not in request.args:
+        return '404'
+    transaction = await Transaction.get(id=request.args.get('id'))
+    return await render_template('receipt.html', transaction=transaction)
+
+
+@app.route('/flag')
+async def flag():
+    return await render_template(
+        'flag.html',
+        balance=session['balance'],
+        flag=os.getenv('FLAG'),
+    )
+
+
+@app.route('/')
+@app.route('/index.html')
+async def index():
+    with open(__file__) as f:
+        return await render_template('source-highlight.html', code=f.read())
+
+
+register_tortoise(
+    app,
+    db_url="sqlite://./data.db",
+    modules={"models": [__name__]},
+    generate_schemas=True,
+)
+
+if __name__ == '__main__':
+    app.run()