diff --git a/dislocker_auth.py b/dislocker_auth.py new file mode 100644 index 0000000..0644e5a --- /dev/null +++ b/dislocker_auth.py @@ -0,0 +1,59 @@ +import psycopg2 +import os +import json +from flask import Flask, request, jsonify + +config_dir_path = "./config/" +db_config_path = config_dir_path + "db.json" +if not os.path.isfile(db_config_path): + if not os.path.isdir(config_dir_path): + os.mkdir(config_dir_path) + + db_config = { + "host": "localhost", + "db": "dislocker", + "username": "user", + "password": "example_pass", + "port": "5432" + } + with open(db_config_path, "w") as w: + json.dump(db_config, w, indent=4) +elif os.path.isfile(db_config_path): + with open(db_config_path, "r") as r: + db_config = json.load(r) + +class Auth(): + def __init__(self, host, db, port, user, password): + self.db = psycopg2.connect(f"host={host} dbname={db} port={port} user={user} password={password}") + + def check(self, pc_number, password): + cursor = self.db.cursor() + cursor.execute("SELECT * FROM pc_list WHERE pc_number = %s AND password_hash = %s", (pc_number, password)) + pc_info = cursor.fetchall() + if not pc_info: + return 1 + else: + return 0 + + def delete(self, pc_number): + cursor = self.db.cursor() + cursor.execute("UPDATE pc_list SET password_hash = NULL WHERE pc_number = %s", (pc_number,)) + self.db.commit() + + +app = Flask(__name__) +auth = Auth(db_config["host"], db_config["db"], db_config["port"], db_config["username"], db_config["password"]) + +@app.route('/verify', methods=['POST']) +def verify(): + pc_number = int(request.json.get('pc_number')) + password = request.json.get('password') + + if auth.check(pc_number, password) == 0: + auth.delete(pc_number) + return jsonify({'message': 'ok'}), 200 + else: + return jsonify({'message': 'damedesu'}), 401 + +if __name__ == '__main__': + app.run() \ No newline at end of file