def register(site):
if request.method == 'POST':
username = request.form['username']
password = request.form['password']
db = get_db()
error = None
if not username:
error = 'Username is required.'
elif not password:
error = 'Password is required.'
elif db.execute('SELECT id FROM user WHERE username = ?',
(username, )).fetchone() is not None:
error = 'User {} is already registered.'.format(username)
if error is None:
db.execute('INSERT INTO user (username, password) VALUES (?, ?)',
(username, generate_password_hash(password)))
db.commit()
user = get_user_from_username(username, db)
db.execute('INSERT INTO balance (user_id, balance) VALUES (?, ?)',
(user['id'], 0))
db.commit()
return redirect(url_for(f'{site.value}.login'))
flash(error)
return render_template(f'register_{site.value}.html')
def get_balance():
user_id = session.get('user_id')
assert user_id is not None
db = get_db()
balance = db.execute('SELECT * FROM balance where user_id = ?',
(user_id, )).fetchone()
return balance['balance']
def load_logged_in_user():
user_id = session.get('user_id')
if user_id is None:
g.user = None
else:
g.user = get_db().execute('SELECT * FROM user WHERE id = ?',
(user_id, )).fetchone()
def update_balance(new_balance):
user_id = session.get('user_id')
assert user_id is not None
db = get_db()
db.execute('UPDATE balance SET balance = ? WHERE user_id = ?',
(new_balance, user_id))
db.commit()
def lookup_user():
user_id = None
form = LookupForm()
if form.validate_on_submit():
username = form.user.data
db = get_db()
user = db.execute(
'SELECT * FROM user WHERE username = ?', (username,)
).fetchone()
if user_id:
user_id = int(user['id'])
elif request.method == 'POST':
flash(form.errors)
return render_template(
'lookup_safe.html', user_id=user_id, form=form)
def lookup_user():
user_id = None
if request.method == 'POST':
username = request.form['user']
db = get_db()
# executescript() allows for multiple commands
user = db.executescript(
f'SELECT * FROM user WHERE username = "******"').fetchone()
# This isn't really needed, but I needed to use executescript(),
# which doesn't seem to work with fetchone(), to make fun SQL
# injections possible.
if not user:
user = db.execute(
f'SELECT * FROM user WHERE username = "******"').fetchone(
)
if user:
user_id = int(user['id'])
return render_template('lookup_vuln.html', user_id=user_id)
def get_user_from_username(username, db=None):
if db is None:
db = get_db()
user = db.execute('SELECT * FROM user WHERE username = ?',
(username, )).fetchone()
return user