0YHkKB8M

1from flask import (
2    Flask,
3    request,
4    render_template_string,
5    session,
6    redirect,
7    send_file
8)
9from random import SystemRandom
10import sqlite3
11import os
12
13app = Flask(__name__)
14app.secret_key = os.getenv('FLASK_KEY')
15
16rand = SystemRandom()
17
18allowed_characters = set(
19    'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ123456789'
20)
21
22
23def execute(query):
24    con = sqlite3.connect('db/db.sqlite3')
25    cur = con.cursor()
26    cur.execute(query)
27    con.commit()
28    return cur.fetchall()
29
30
31def generate_token():
32    return ''.join(
33        rand.choice(list(allowed_characters)) for _ in range(32)
34    )
35
36
37def create_user(username, password):
38    if any(c not in allowed_characters for c in username):
39        return (False, 'Alphanumeric usernames only, please.')
40    if len(username) < 1:
41        return (False, 'Username is too short.')
42    if len(password) > 50:
43        return (False, 'Password is too long.')
44    other_users = execute(
45        f'SELECT * FROM users WHERE username=\'{username}\';'
46    )
47    if len(other_users) > 0:
48        return (False, 'Username taken.')
49    execute(
50        'INSERT INTO users (username, password)'
51        f'VALUES (\'{username}\', \'{password}\');'
52    )
53    return (True, '')
54
55
56def check_login(username, password):
57    if any(c not in allowed_characters for c in username):
58        return False
59    correct_password = execute(
60        f'SELECT password FROM users WHERE username=\'{username}\';'
61    )
62    if len(correct_password) < 1:
63        return False
64    return correct_password[0][0] == password
65
66
67@app.route('/', methods=['GET', 'POST'])
68def login():
69    error = ''
70    if request.method == 'POST':
71        valid_login = check_login(
72            request.form['username'],
73            request.form['password']
74        )
75        if valid_login:
76            session['username'] = request.form['username']
77            return redirect('/message')
78        error = 'Incorrect username or password.'
79    if 'username' in session:
80        return redirect('/message')
81    return render_template_string('''
82        <link rel="stylesheet" href="/static/style.css" />
83        <div class="container">
84            <p>Log in to see Aaron's message!</p>
85            <form method="POST">
86                <label for="username">Username</label>
87                <input type="text" name="username" />
88                <label for="password">Password</label>
89                <input type="password" name="password" />
90                <input type="submit" value="Log In" />
91            </form>
92            <p>{{ error }}</p>
93            <a href="/register">Register</a>
94        <div class="container">
95    ''', error=error)
96
97
98@app.route('/register', methods=['GET', 'POST'])
99def register():
100    message = ''
101    if request.method == 'POST':
102        success, message = create_user(
103            request.form['username'],
104            request.form['password']
105        )
106        if success:
107            session['username'] = request.form['username']
108            return redirect('/message')
109    return render_template_string('''
110        <link rel="stylesheet" href="/static/style.css" />
111        <div class="container">
112            <p>Register!</p>
113            <form method="POST">
114                <label for="username">Username</label>
115                <input type="text" name="username" />
116                <label for="password">Password</label>
117                <input type="password" name="password" />
118                <input type="submit" value="Register" />
119            </form>
120            <p>{{ error }}</p>
121        </div>
122    ''', error=message)
123
124
125@app.route('/message')
126def message():
127    if 'username' not in session:
128        return redirect('/')
129    if session['username'] == 'ginkoid':
130        return send_file(
131            'flag.mp3',
132            attachment_filename='flag-at-end-of-file.mp3'
133        )
134    return '''
135        <link rel="stylesheet" href="/static/style.css" />
136            <div class="container">
137            <p>You are logged in!</p>
138            <p>Unfortunately, Aaron's message is for cool people only.</p>
139            <p>(like ginkoid)</p>
140            <a href="/logout">Log out</a>
141        </div>
142    '''
143
144
145@app.route('/logout')
146def logout():
147    if 'username' not in session:
148        return redirect('/')
149    del session['username']
150    return redirect('/')
151
152
153def init():
154    # this is terrible but who cares
155    execute('''
156        CREATE TABLE IF NOT EXISTS users (
157            username TEXT PRIMARY KEY,
158            password TEXT
159        );
160    ''')
161    execute('DROP TABLE users;')
162    execute('''
163        CREATE TABLE users (
164            username TEXT PRIMARY KEY,
165            password TEXT
166        );
167    ''')
168
169    # put ginkoid into db
170    ginkoid_password = generate_token()
171    execute(
172        'INSERT OR IGNORE INTO users (username, password)'
173        f'VALUES (\'ginkoid\', \'{ginkoid_password}\');'
174    )
175    execute(
176        f'UPDATE users SET password=\'{ginkoid_password}\''
177        f'WHERE username=\'ginkoid\';'
178    )
179
180
181init()
182