114 lines
3.3 KiB
Python
114 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
from burvesigner import BurveSigner
|
|
from flask import Flask, make_response, render_template, request, flash
|
|
from hashlib import sha256
|
|
from time import time
|
|
import base64
|
|
import json
|
|
import secrets
|
|
import sqlite3
|
|
|
|
|
|
app = Flask(__name__)
|
|
app.config["SECRET_KEY"] = secrets.token_hex()
|
|
db = "database.db"
|
|
signer = BurveSigner()
|
|
|
|
|
|
b64p = lambda x: x + b"=" * (-len(x) % 4)
|
|
b64u = lambda x: x.rstrip(b"=")
|
|
b64e = lambda x: b64u(base64.urlsafe_b64encode(x))
|
|
b64d = lambda x: base64.urlsafe_b64decode(b64p(x))
|
|
|
|
|
|
def generate_token(user, role):
|
|
payload = b64e(json.dumps({
|
|
"user": user,
|
|
"role": role,
|
|
"exp": int(time()) + 300
|
|
}).encode())
|
|
signature = signer.sign(payload)
|
|
token = payload + b"." + signature
|
|
return token.decode()
|
|
|
|
|
|
def validate_token(token):
|
|
assert token.count(".") == 1
|
|
payload, signature = token.split(".")
|
|
assert signer.verify(payload.encode(), signature.encode())
|
|
payload = json.loads(b64d(payload.encode()).decode())
|
|
assert payload.get("user") and payload.get("role") and payload.get("exp")
|
|
user, role, exp = payload.get("user"), payload.get("role"), payload.get("exp")
|
|
assert int(time()) < exp
|
|
return [user, role, exp]
|
|
|
|
|
|
def create_response_data(user, role):
|
|
data = f"Welcome, {user}! "
|
|
if role == "admin":
|
|
try:
|
|
data += open("/flag.txt").read()
|
|
except:
|
|
data += "Flag is missing, please contact probset"
|
|
else:
|
|
data += "Only admin can access the flag"
|
|
return data
|
|
|
|
|
|
def get_role(username, password):
|
|
connection = sqlite3.connect(db)
|
|
cursor = connection.cursor()
|
|
cursor.execute(
|
|
"SELECT role FROM users WHERE username = ? AND password = ?",
|
|
(username, sha256(b"gema" + password.encode() + b"steg").hexdigest()),
|
|
)
|
|
role = cursor.fetchone()
|
|
connection.close()
|
|
assert role
|
|
return role[0]
|
|
|
|
|
|
@app.route("/params", methods=["GET"])
|
|
def params():
|
|
return f"<pre>{json.dumps(signer.get_params(), indent=2)}</pre>"
|
|
|
|
|
|
@app.route("/", methods=["POST", "GET"])
|
|
def index():
|
|
if request.method == "POST":
|
|
username = request.form.get("username")
|
|
password = request.form.get("password")
|
|
try:
|
|
assert 2 <= len(username) <= 32
|
|
assert 2 <= len(password) <= 32
|
|
role = get_role(username, password)
|
|
token = generate_token(username, role)
|
|
validate_token(token)
|
|
response = make_response(render_template(
|
|
"index.html",
|
|
data=create_response_data(username, role)
|
|
))
|
|
response.set_cookie("token", token)
|
|
return response
|
|
except:
|
|
flash("Invalid username and/or password")
|
|
return render_template("index.html")
|
|
else:
|
|
token = request.cookies.get("token")
|
|
if token:
|
|
try:
|
|
username, role, _ = validate_token(token)
|
|
return render_template(
|
|
"index.html",
|
|
data=create_response_data(username, role)
|
|
)
|
|
except:
|
|
flash("Invalid token")
|
|
return render_template("index.html")
|
|
else:
|
|
return render_template("index.html")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host="0.0.0.0", port=80, debug=False)
|