from flask import Flask, g, request, render_template, redirect, current_app, send_file
from flask_argon2 import Argon2
from io import BytesIO
import jwt
import time
import sqlite3
import secrets

app = Flask(__name__)
app.config['PRIVATE_KEY'] = '''
-----BEGIN RSA PRIVATE KEY-----
MIICXAIBAAKBgQCo0oR5UXusQs7tvgcT5EOLB+2JaKXmmip3ViGijLHku2Y+gyas
0KFYHEQjTgz2AH1N9sUu8tzLAxn+wun28qyF3Paswmx27J/pAmbX8v6g+oGur4Jz
V6ZoM5PA5iD5UvWYcLBczB84GcqhQkLHh8n/sZXP9jXMnxjTPD4nuPuQ3wIDAQAB
AoGAd0s6/RdNEu6qlmifS7kS2V2ixmRCRu9NbsJYRiqxUfXyS94VKCzMthxTMbdn
hTXXVY44y/IlfvcUGWfWOABHU7JK5NfWbwJfH0dU2kNEf8LzPmf1DzGy6vj01i/u
6KrSJMqJOW62NxQ1GjkvWVGgoy8RrHKzrkM7bnQ+i6JDVLECQQDX24V58LQwMmo2
JDZHjEZLZlx4xQz3lzhrLOfn7B3zgspRrgufOp2SaL+nFaphUl4w8P9mo/FcDmqc
R402Z4yTAkEAyDfCiZiGBgPm7mDOLiJ1Wpyc21fsF6zwoc56xSbeK+a3t9LxT00M
1W+qZv6e89erUmGNl85CwFmoyMEPdIgmBQJAOOUZp2x0cge3yxF8ZRtqI9GVKhf2
NQRc0JMDhTPNKTQeE61mTs/qXH7TlTy2rfRB83ByQSGRKox6OTr6044zlQJAKuCe
Hb93PESLqRM8NG8WuL//a43ptqxHoC9K5XvMapRvVcOr//KdQ/w0/veabNgMDYls
vEzkyLKqzctilu8tTQJBAMXgHKX62GhH54pB4GrLLS25JpxqUNChDuPGaMPilfCW
WOsFVC93MPtLA/YaAJHKZNoaXulkb5q3jhlWxCpDAKM=
-----END RSA PRIVATE KEY-----
'''

app.config['PUBLIC_KEY'] = '''-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCo0oR5UXusQs7tvgcT5EOLB+2J
aKXmmip3ViGijLHku2Y+gyas0KFYHEQjTgz2AH1N9sUu8tzLAxn+wun28qyF3Pas
wmx27J/pAmbX8v6g+oGur4JzV6ZoM5PA5iD5UvWYcLBczB84GcqhQkLHh8n/sZXP
9jXMnxjTPD4nuPuQ3wIDAQAB
-----END PUBLIC KEY-----'''

argon2 = Argon2(app)

flag1 = 'HM{sql_inj3ctions_4re_still_r3l3vant}'
flag2 = 'HM{y0u_th0ught_p4ssword_ha5h1ng_is_enough?}'
flag3 = 'HM{s3ssion_manag3ment_is_super_h4rd}'

database = 'evilcorp.sqlite3'

def check_session(cookies):

    if 'session' not in cookies:
        return False

    session = cookies['session']
    try:
        session_decoded = jwt.decode(session, app.config['PUBLIC_KEY'])
    except jwt.InvalidTokenError:
        return False
    return session_decoded

def get_db():
    db = getattr(g, '_database', None)
    if db is None:
        db = g._database = sqlite3.connect(database)
    return db


@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()

@app.route('/', methods=['GET'])
def show_index():
    session = check_session(request.cookies)
    if session:
        authenticated = True
        username = session['user']
    else:
        authenticated = False
        username = None
    return render_template('index.html', authenticated=authenticated, username=username)

@app.route('/news', methods=['GET'])
def show_news():
    session = check_session(request.cookies)
    if session:
        authenticated = True
        username = session['user']
    else:
        authenticated = False
        username = None

    try:
        cur = get_db().execute('SELECT uid, title, body, images FROM news ORDER BY uid DESC')
        news = cur.fetchall()
    except sqlite3.Error as e:
        print(e)
        return 'Something went wrong, ping the admins'
    
    return render_template('news.html', news=news, authenticated=authenticated, username=username)

@app.route('/images/<uid>', methods=['GET'])
def show_image(uid):
    try:
        cur = get_db().execute('SELECT uid, name, body FROM images WHERE uid = ' + uid)
        image = cur.fetchone()
    except sqlite3.Error as e:
        print(e)
        return 'Something went wrong, ping the admins'
    if not image:
    	return 'No image found'
    return send_file(BytesIO(image[2]), attachment_filename=image[1]+'.jpg', mimetype='image/jpg')

@app.route('/login', methods=['GET'])
def show_login():
    return render_template('login.html')

@app.route('/login', methods=['POST'])
def login():

    if 'username' not in request.form or 'password' not in request.form:
        return 'All fields are required!'
    username = request.form['username']
    password = request.form['password']

    #if not username.isalnum():
    #    return 'Login failed'

    if len(password) < 6:
        return 'Login failed'

    try:
        cur = get_db().execute('SELECT uid, username, password FROM users WHERE username = ? LIMIT 1', (username,))
        cur_user = cur.fetchone()
    except sqlite3.Error as e:
        print(e)
        return 'Something went wrong, ping the admins'

    if not cur_user:
        return 'Login failed'

    if not argon2.check_password_hash(cur_user[2], password):
        return 'Login failed'

    session = jwt.encode({'user': cur_user[1], 'admin': False, 'iat': int(time.time())}, key=app.config['PRIVATE_KEY'] , algorithm='RS256')

    response = current_app.make_response(redirect('/user', 302))
    response.set_cookie('session', value=session, path='/', httponly=True)
    return response


@app.route('/reset', methods=['GET'])
def show_reset():
    return render_template('reset.html')


@app.route('/reset', methods=['POST'])
def reset():
    if 'email' not in request.form:
        return 'Email is required required!'
    email = request.form['email']

    if len(email) < 6:
        return 'Email too short'

    try:
        cur = get_db().execute('SELECT uid, email, username, recovered, admin FROM users WHERE email = ? LIMIT 1', (email,))
        cur_user = cur.fetchone()
    except sqlite3.Error as e:
        print(e)
        return 'Something went wrong, ping the admins'

    if not cur_user:
        return 'No user found'

    if cur_user[4]:
        return 'Nope'

    if cur_user[3]:
        return 'Password for this user has already been reset'

    uid = cur_user[0]
    token = secrets.token_hex()

    try:
        #cur = get_db().execute('UPDATE user SET recovery = ?, recovered = 1 WHERE email = ?', (token, email,))
        cur = get_db().execute('UPDATE users SET recovery = ? WHERE uid = ?', (token, uid,))
        get_db().commit()
    except sqlite3.Error as e:
        print(e)
        return 'Something went wrong, ping the admins'

    return 'Reset token sent to user email'

@app.route('/reset2', methods=['POST'])
def reset2():
    if 'email' not in request.form or 'token' not in request.form:
        return 'Both fields are required!'

    email = request.form['email']
    token = request.form['token']

    try:
        cur = get_db().execute('SELECT * FROM users WHERE recovery = ? AND email = ? AND recovered != 1', (token, email,))
        cur_user = cur.fetchone()
    except sqlite3.Error as e:
        print(e)
        return 'Something went wrong, ping the admins'

    if not cur_user:
        return 'Wrong email or token or password already reset one time'

    else:
        newpassword = secrets.token_urlsafe(16)
        try:
            cur = get_db().execute('UPDATE users SET password = ?, recovered = 1 WHERE recovery = ? AND email = ?', (argon2.generate_password_hash(newpassword), token, email,))
            get_db().commit()
        except sqlite3.Error as e:
            print(e)
            return 'Something went wrong, ping the admins'

    return 'Your new password is ' + newpassword


@app.route('/user', methods=['GET'])
def show_user():
    session = check_session(request.cookies)
    if not session:
        return redirect('/login', 302)
    else:
        if not session['admin']:
            return render_template('user.html', authenticated=True, username=session['user'], flag2=flag2)
        return 'Good job! Here\'s the last flag: ' + flag3

@app.route('/logout', methods=['GET'])
def logout():
    response = current_app.make_response(redirect('/', 302))
    for cookie in request.cookies:
        response.set_cookie(cookie, value='', expires=0)
    return response

@app.route('/robots.txt', methods=['GET'])
def show_robots():
    return 'Disallow: /pub.asc'

@app.route('/pub.asc', methods=['GET'])
def show_pubkey():
    return app.config['PUBLIC_KEY']

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=7000)