From 6bc3ba9ad7cabffe9b29f501e88c9576d0974868 Mon Sep 17 00:00:00 2001 From: Mateusz779 Date: Tue, 18 Apr 2023 14:51:04 +0200 Subject: [PATCH] code cleanup --- db.py | 190 ++++++++++++++++------------------------------------------ 1 file changed, 50 insertions(+), 140 deletions(-) diff --git a/db.py b/db.py index 1b57d6f..e0e6695 100644 --- a/db.py +++ b/db.py @@ -2,6 +2,7 @@ import psycopg2 import config import utils + def connect(): global cur, conn try: @@ -12,9 +13,9 @@ def connect(): port=config.port) except Exception as ex: print(f"Error connecting to PostgreSQL: {ex}") - + cur = conn.cursor() - + with conn.cursor() as cur: cur.execute(""" CREATE TABLE IF NOT EXISTS image ( @@ -47,120 +48,77 @@ def connect(): client_ip INET );""") conn.commit() - + + def get_cur(): return conn.cursor() + def get_conn(): return conn + def get_one(sql, value): connect() with get_cur() as cur: - cur.execute(sql,(value,)) + cur.execute(sql, (value,)) try: return cur.fetchone()[0] except: return None + def add_conf_image(name, token): connect() with get_cur() as cur: cur.execute(""" INSERT INTO image (image_name, token) VALUES (%s, %s) - """,(name, token,)) + """, (name, token,)) conn.commit() - -# def get_conf(sql,value): -# connect() -# with get_cur() as cur: -# cur.execute(sql,(value,)) -# try: -# return cur.fetchone()[0] -# except: -# return None + def get_conf_image(token): return get_one("SELECT image_name FROM image WHERE token = %s", 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_one("SELECT id FROM image WHERE token = %s", 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): return get_one("SELECT id FROM image WHERE image_name = %s", 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() with get_cur() as cur: cur.execute(""" INSERT INTO users (username, password) VALUES (%s, %s) - """,(username, utils.hash_password(password),)) + """, (username, utils.hash_password(password),)) conn.commit() - + + def get_user_pass(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),)) + """, (username, utils.hash_password(password),)) try: return cur.fetchone()[0] except: return None + def get_user_byid(id): return get_one("SELECT id FROM users WHERE id = %s", 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): return get_one("SELECT user_id FROM auth_tokens WHERE token = %s", 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() @@ -169,18 +127,19 @@ def add_auth_token(user_id): cur.execute(""" INSERT INTO auth_tokens (user_id, token, expires_on) VALUES (%s, %s, NOW() + INTERVAL '1 day') - """,(user_id,token,)) + """, (user_id, token,)) conn.commit() return token - + + def login(username, password): user_id = get_user_pass(username, password) if user_id is not None: return add_auth_token(user_id) else: return None - - + + def get_image_allocation_all(): connect() with get_cur() as cur: @@ -202,85 +161,53 @@ def get_image_allocation(sql, value): except: return None + def get_image_allocation_time(token): image_id = get_conf_id(token) if image_id is None: return None return get_one("SELECT last_access_time FROM image_allocation WHERE image_id = %s", image_id) -# def get_image_allocation_time_imageid(image_id): -# get_image_allocation("SELECT last_access_time FROM image_allocation WHERE image_id = %s", 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): get_one("SELECT last_access_time FROM image_allocation WHERE id = %s", 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 - + return get_one("SELECT last_access_time FROM image_allocation WHERE id = %s", 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): - return get_one("SELECT client_ip FROM image_allocation WHERE id = %s", id) - # connect() - # with get_cur() as cur: - # cur.execute(""" - # SELECT client_ip FROM image_allocation WHERE id = %s - # """,(id,)) - # try: - # return cur.fetchone()[0] - # except: - # return None + return get_one("SELECT client_ip FROM image_allocation WHERE id = %s", id) + def set_image_allocation(token, client_ip): id_image = get_conf_id(token) if id_image is None: return None - + connect() with get_cur() as cur: cur.execute(""" INSERT INTO image_allocation (image_id, client_ip, last_access_time) VALUES (%s, %s, NOW()) - """,(id_image,client_ip,)) + """, (id_image, client_ip,)) conn.commit() return token + def del_image_allocation_token(token): id_image = get_conf_id(token) if id_image is None: - return None - + return None + return del_image_allocation_id_image(id_image) + def del_image_allocation(sql, value): connect() with get_cur() as cur: @@ -291,40 +218,23 @@ def del_image_allocation(sql, value): except: return None + def del_image_allocation_id_image(image_id): return del_image_allocation("DELETE FROM image_allocation WHERE image_id = %s", 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("DELETE FROM image_allocation WHERE id = %s", 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): + +def del_image_allocation_id(id): + return del_image_allocation("DELETE FROM image_allocation WHERE id = %s", id) + + +def update_image_allocation_time(id): connect() with get_cur() as cur: cur.execute(""" UPDATE image_allocation SET last_access_time = NOW() WHERE id = %s - """,(id,)) + """, (id,)) try: conn.commit() return True except: - return None \ No newline at end of file + return None