From 3edab063283c4b01b3e0f69223acd3c2836d836c Mon Sep 17 00:00:00 2001 From: Mateusz779 Date: Tue, 18 Apr 2023 14:11:48 +0200 Subject: [PATCH] code cleanup --- app.py | 2 +- db.py | 185 ++++++++++++++++++++++++++++++++++----------------------- 2 files changed, 111 insertions(+), 76 deletions(-) diff --git a/app.py b/app.py index 8deffd1..2fb76ee 100644 --- a/app.py +++ b/app.py @@ -177,6 +177,6 @@ def add_ip(): except: pass print(token, ip) - return + return "" if __name__ == '__main__': app.run(host="0.0.0.0") \ No newline at end of file diff --git a/db.py b/db.py index 7400161..33ef699 100644 --- a/db.py +++ b/db.py @@ -11,7 +11,7 @@ def connect(): password=config.password, port=config.port) except Exception as ex: - print(f"Error connecting to PostgreSQL: {e}") + print(f"Error connecting to PostgreSQL: {ex}") cur = conn.cursor() @@ -63,38 +63,52 @@ def add_conf_image(name, token): """,(name, token,)) conn.commit() -def get_conf_image(token): +def get_conf(columns, param, value): connect() with get_cur() as cur: cur.execute(""" - SELECT image_name FROM image WHERE token = %s - """,(token,)) - try: - return cur.fetchone()[0] - except: - return None - -def get_conf_id(token): - connect() - with get_cur() as cur: - cur.execute(""" - SELECT id FROM image WHERE token = %s - """,(token,)) + SELECT %s FROM image WHERE %s = %s + """,(columns, param, value,)) try: return cur.fetchone()[0] except: return None +def get_conf_image(token): + return get_conf("image_name", "token", token) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # SELECT image_name FROM image WHERE token = %s + # """,(token,)) + # try: + # return cur.fetchone()[0] + # except: + # return None + +def get_conf_id(token): + return get_conf("id", "token", token) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # SELECT id FROM image WHERE token = %s + # """,(token,)) + # try: + # return cur.fetchone()[0] + # except: + # return None + def get_conf_id_name(name): - connect() - with get_cur() as cur: - cur.execute(""" - SELECT id FROM image WHERE image_name = %s - """,(name+".squashfs",)) - try: - return cur.fetchone()[0] - except: - return None + return get_conf("id", "image_name", name) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # SELECT id FROM image WHERE image_name = %s + # """,(name+".squashfs",)) + # try: + # return cur.fetchone()[0] + # except: + # return None def add_user(username, password): connect() @@ -105,7 +119,7 @@ def add_user(username, password): """,(username, utils.hash_password(password),)) conn.commit() -def get_user(username, password): +def get_user_pass(username, password): connect() with get_cur() as cur: cur.execute(""" @@ -115,7 +129,7 @@ def get_user(username, password): return cur.fetchone()[0] except: return None - + def get_user_byid(id): connect() with get_cur() as cur: @@ -150,7 +164,7 @@ def add_auth_token(user_id): return token def login(username, password): - user_id = get_user(username, password) + user_id = get_user_pass(username, password) if user_id is not None: return add_auth_token(user_id) else: @@ -168,57 +182,64 @@ def get_image_allocation_all(): except: return None -def get_image_allocation_time(token): - id_image = get_conf_id(token) - if id_image is None: - return None - + +def get_image_allocation(column, param, value): connect() with get_cur() as cur: cur.execute(""" - SELECT last_access_time FROM image_allocation WHERE image_id = %s - """,(id_image,)) + SELECT %s FROM image_allocation WHERE %s = %s + """,(column, param, value,)) try: return cur.fetchone()[0] except: return None +def get_image_allocation_time(token): + image_id = get_conf_id(token) + if image_id is None: + return None + + return get_image_allocation_time_imageid(image_id) + def get_image_allocation_time_imageid(image_id): - connect() - with get_cur() as cur: - cur.execute(""" - SELECT last_access_time FROM image_allocation WHERE image_id = %s - """,(image_id,)) - try: - return cur.fetchone()[0] - except: - return None + get_image_allocation("last_access_time", "image_id", image_id) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # SELECT last_access_time FROM image_allocation WHERE image_id = %s + # """,(image_id,)) + # try: + # return cur.fetchone()[0] + # except: + # return None def get_image_allocation_time_id(id): - connect() - with get_cur() as cur: - cur.execute(""" - SELECT last_access_time FROM image_allocation WHERE id = %s - """,(id,)) - try: - return cur.fetchone()[0] - except: - return None + get_image_allocation("last_access_time", "id", id) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # SELECT last_access_time FROM image_allocation WHERE id = %s + # """,(id,)) + # try: + # return cur.fetchone()[0] + # except: + # return None def get_image_allocation_clientip(token): id_image = get_conf_id(token) if id_image is None: return None - connect() - with get_cur() as cur: - cur.execute(""" - SELECT client_ip FROM image_allocation WHERE image_id = %s - """,(id_image,)) - try: - return cur.fetchone()[0] - except: - return None + get_image_allocation("client_ip", "image_id", id_image) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # SELECT client_ip FROM image_allocation WHERE image_id = %s + # """,(id_image,)) + # try: + # return cur.fetchone()[0] + # except: + # return None def get_image_allocation_clientip_id(id): connect() @@ -252,29 +273,43 @@ def del_image_allocation_token(token): return del_image_allocation_id_image(id_image) -def del_image_allocation_id_image(id_image): +def del_image_allocation(column, value): connect() with get_cur() as cur: cur.execute(""" - DELETE FROM image_allocation WHERE image_id = %s - """,(id_image,)) + DELETE FROM image_allocation WHERE %s = %s + """,(column,value, )) try: conn.commit() return True except: return None -def del_image_allocation_id(id): - connect() - with get_cur() as cur: - cur.execute(""" - DELETE FROM image_allocation WHERE id = %s - """,(id,)) - try: - conn.commit() - return True - except: - return None +def del_image_allocation_id_image(image_id): + return del_image_allocation("image_id", image_id) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # DELETE FROM image_allocation WHERE image_id = %s + # """,(id_image,)) + # try: + # conn.commit() + # return True + # except: + # return None + +def del_image_allocation_id(id): + return del_image_allocation("id", id) + # connect() + # with get_cur() as cur: + # cur.execute(""" + # DELETE FROM image_allocation WHERE id = %s + # """,(id,)) + # try: + # conn.commit() + # return True + # except: + # return None def update_image_allocation_time(id): connect()