diff --git a/app.py b/app.py index 059b121..2de56a1 100644 --- a/app.py +++ b/app.py @@ -1,4 +1,4 @@ -from flask import Flask, send_file, jsonify, request +from flask import Flask, flash, make_response, redirect, send_file, jsonify, request, render_template, url_for import db import os from werkzeug.utils import secure_filename @@ -15,6 +15,30 @@ def ssh_thread_function(): ssh_thread = threading.Thread(target=ssh_thread_function) ssh_thread.start() +@app.route('/') +def main(): + auth_token = request.cookies.get('auth_token') + if auth_token != "" and auth_token is not None: + if db.get_user_bytoken(auth_token) is None: + return render_template('template/login.html') + return render_template('template/index.html') + +@app.route('/api/login', methods=['POST']) +def login_post(): + username = request.form['username'] + password = request.form['password'] + + auth_token = db.login() + if auth_token is None: + flash('Nieprawidłowa nazwa użytkownika lub hasło.', 'error') + return redirect(url_for('login')) + + response = make_response(render_template('template/index.html')) + response.set_cookie('auth_token', auth_token) + + return response + + @app.route("/api/addimage", methods=['POST']) def add_image(): db.Connect() @@ -51,19 +75,14 @@ def add_image(): else: filename = "1"+filename file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename)) - - db.AddVPNImage(filename, token) + db.add_conf_image(filename, token) return jsonify(message="ok") -@app.route("/api/getvpn") +@app.route("/api/getconf") def get_image(): - try: - filename = db.GetVPNImage(request.headers['token'])[0] - print(filename) - except: - filename = "default.squashfs" + filename = db.get_conf_image(request.headers['token']) if filename is None or filename == "": filename = "default.squashfs" diff --git a/db.py b/db.py index ace7417..9f26f92 100644 --- a/db.py +++ b/db.py @@ -1,7 +1,8 @@ import psycopg2 import config +import utils -def Connect(): +def connect(): global cur, conn try: conn = psycopg2.connect(database=config.database, @@ -21,31 +22,107 @@ def Connect(): image_name VARCHAR(255) NOT NULL, token VARCHAR(255) NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP - ) - """) + );""") + cur.execute(""" + CREATE TABLE IF NOT EXISTS users ( + id SERIAL PRIMARY KEY, + username VARCHAR(50) UNIQUE NOT NULL, + password VARCHAR(256) NOT NULL, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP + );""") + cur.execute(""" + CREATE TABLE auth_tokens ( + id SERIAL PRIMARY KEY, + user_id INTEGER NOT NULL REFERENCES users(id), + token VARCHAR(64) NOT NULL, + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP + expires_on TIMESTAMP NOT NULL + );""") conn.commit() -def GetCur(): +def get_cur(): return conn.cursor() -def GetConn(): +def get_conn(): return conn -def AddVPNImage(name, token): - Connect() - with GetCur() as cur: +def add_conf_image(name, token): + connect() + with get_cur() as cur: cur.execute(""" INSERT INTO vpn (image_name, token) VALUES (%s, %s) """,(name, token,)) conn.commit() -def GetVPNImage(token): - Connect() - with GetCur() as cur: +def get_conf_image(token): + connect() + with get_cur() as cur: cur.execute(""" SELECT image_name FROM vpn WHERE token = %s """,(token,)) - return cur.fetchone() + try: + return cur.fetchone()[0] + except: + return None + +def add_user(username, password): + connect() + with get_cur() as cur: + cur.execute(""" + INSERT INTO users (username, password) + VALUES (%s, %s) + """,(username, utils.hash_password(password),)) + conn.commit() - \ No newline at end of file +def get_user(username, password): + connect() + with get_cur() as cur: + cur.execute(""" + SELECT id FROM users WHERE username = %s AND password = %s + """,(username, utils.hash_password(password),)) + try: + return cur.fetchone()[0] + except: + return None + +def get_user_byid(id): + connect() + with get_cur() as cur: + cur.execute(""" + SELECT id FROM users WHERE id = %s + """,(id,)) + try: + return cur.fetchone()[0] + except: + return None + +def get_user_bytoken(token): + connect() + with get_cur() as cur: + cur.execute(""" + SELECT user_id FROM auth_tokens WHERE token = %s + """,(token,)) + try: + return cur.fetchone()[0] + except: + return None + +def add_auth_token(user_id): + token = utils.generate_auth_token() + connect() + with get_cur() as cur: + cur.execute(""" + INSERT INTO auth_tokens (user_id, token) + VALUES (%s, %s) + """,(user_id,token,)) + conn.commit() + return token + +def login(username, password): + user_id = get_user(username, password) + if user_id is not None: + return add_auth_token(user_id) + else: + return None + diff --git a/template/login.html b/template/login.html new file mode 100644 index 0000000..b941610 --- /dev/null +++ b/template/login.html @@ -0,0 +1,63 @@ + + + + + Strona logowania + + + +
+

Logowanie

+
+ + + + + + + +
+
+ + diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..75f19f5 --- /dev/null +++ b/utils.py @@ -0,0 +1,8 @@ +import hashlib +import secrets + +def hash_password(password): + return hashlib.sha512(password.encode('utf-8')).hexdigest() + +def generate_auth_token(): + return secrets.token_urlsafe(32) \ No newline at end of file