code cleanup

This commit is contained in:
Mateusz779 2023-04-18 14:11:48 +02:00
parent 0b6c4197c9
commit 3edab06328
2 changed files with 111 additions and 76 deletions

2
app.py
View File

@ -177,6 +177,6 @@ def add_ip():
except: except:
pass pass
print(token, ip) print(token, ip)
return return ""
if __name__ == '__main__': if __name__ == '__main__':
app.run(host="0.0.0.0") app.run(host="0.0.0.0")

175
db.py
View File

@ -11,7 +11,7 @@ def connect():
password=config.password, password=config.password,
port=config.port) port=config.port)
except Exception as ex: except Exception as ex:
print(f"Error connecting to PostgreSQL: {e}") print(f"Error connecting to PostgreSQL: {ex}")
cur = conn.cursor() cur = conn.cursor()
@ -63,38 +63,52 @@ def add_conf_image(name, token):
""",(name, token,)) """,(name, token,))
conn.commit() conn.commit()
def get_conf_image(token): def get_conf(columns, param, value):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
SELECT image_name FROM image WHERE token = %s SELECT %s FROM image WHERE %s = %s
""",(token,)) """,(columns, param, value,))
try: try:
return cur.fetchone()[0] return cur.fetchone()[0]
except: except:
return None 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): def get_conf_id(token):
connect() return get_conf("id", "token", token)
with get_cur() as cur: # connect()
cur.execute(""" # with get_cur() as cur:
SELECT id FROM image WHERE token = %s # cur.execute("""
""",(token,)) # SELECT id FROM image WHERE token = %s
try: # """,(token,))
return cur.fetchone()[0] # try:
except: # return cur.fetchone()[0]
return None # except:
# return None
def get_conf_id_name(name): def get_conf_id_name(name):
connect() return get_conf("id", "image_name", name)
with get_cur() as cur: # connect()
cur.execute(""" # with get_cur() as cur:
SELECT id FROM image WHERE image_name = %s # cur.execute("""
""",(name+".squashfs",)) # SELECT id FROM image WHERE image_name = %s
try: # """,(name+".squashfs",))
return cur.fetchone()[0] # try:
except: # return cur.fetchone()[0]
return None # except:
# return None
def add_user(username, password): def add_user(username, password):
connect() connect()
@ -105,7 +119,7 @@ def add_user(username, password):
""",(username, utils.hash_password(password),)) """,(username, utils.hash_password(password),))
conn.commit() conn.commit()
def get_user(username, password): def get_user_pass(username, password):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
@ -150,7 +164,7 @@ def add_auth_token(user_id):
return token return token
def login(username, password): def login(username, password):
user_id = get_user(username, password) user_id = get_user_pass(username, password)
if user_id is not None: if user_id is not None:
return add_auth_token(user_id) return add_auth_token(user_id)
else: else:
@ -168,57 +182,64 @@ def get_image_allocation_all():
except: except:
return None 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() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
SELECT last_access_time FROM image_allocation WHERE image_id = %s SELECT %s FROM image_allocation WHERE %s = %s
""",(id_image,)) """,(column, param, value,))
try: try:
return cur.fetchone()[0] return cur.fetchone()[0]
except: except:
return None 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): def get_image_allocation_time_imageid(image_id):
connect() get_image_allocation("last_access_time", "image_id", image_id)
with get_cur() as cur: # connect()
cur.execute(""" # with get_cur() as cur:
SELECT last_access_time FROM image_allocation WHERE image_id = %s # cur.execute("""
""",(image_id,)) # SELECT last_access_time FROM image_allocation WHERE image_id = %s
try: # """,(image_id,))
return cur.fetchone()[0] # try:
except: # return cur.fetchone()[0]
return None # except:
# return None
def get_image_allocation_time_id(id): def get_image_allocation_time_id(id):
connect() get_image_allocation("last_access_time", "id", id)
with get_cur() as cur: # connect()
cur.execute(""" # with get_cur() as cur:
SELECT last_access_time FROM image_allocation WHERE id = %s # cur.execute("""
""",(id,)) # SELECT last_access_time FROM image_allocation WHERE id = %s
try: # """,(id,))
return cur.fetchone()[0] # try:
except: # return cur.fetchone()[0]
return None # except:
# return None
def get_image_allocation_clientip(token): def get_image_allocation_clientip(token):
id_image = get_conf_id(token) id_image = get_conf_id(token)
if id_image is None: if id_image is None:
return None return None
connect() get_image_allocation("client_ip", "image_id", id_image)
with get_cur() as cur: # connect()
cur.execute(""" # with get_cur() as cur:
SELECT client_ip FROM image_allocation WHERE image_id = %s # cur.execute("""
""",(id_image,)) # SELECT client_ip FROM image_allocation WHERE image_id = %s
try: # """,(id_image,))
return cur.fetchone()[0] # try:
except: # return cur.fetchone()[0]
return None # except:
# return None
def get_image_allocation_clientip_id(id): def get_image_allocation_clientip_id(id):
connect() connect()
@ -252,29 +273,43 @@ def del_image_allocation_token(token):
return del_image_allocation_id_image(id_image) return del_image_allocation_id_image(id_image)
def del_image_allocation_id_image(id_image): def del_image_allocation(column, value):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
DELETE FROM image_allocation WHERE image_id = %s DELETE FROM image_allocation WHERE %s = %s
""",(id_image,)) """,(column,value, ))
try: try:
conn.commit() conn.commit()
return True return True
except: except:
return None 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): def del_image_allocation_id(id):
connect() return del_image_allocation("id", id)
with get_cur() as cur: # connect()
cur.execute(""" # with get_cur() as cur:
DELETE FROM image_allocation WHERE id = %s # cur.execute("""
""",(id,)) # DELETE FROM image_allocation WHERE id = %s
try: # """,(id,))
conn.commit() # try:
return True # conn.commit()
except: # return True
return None # except:
# return None
def update_image_allocation_time(id): def update_image_allocation_time(id):
connect() connect()