public/Zdalne-systemy-z-kit-uVPN
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

code cleanup

Browse Source
This commit is contained in:
Mateusz779 2023-04-18 14:51:04 +02:00
parent b5d0cc390a
commit 6bc3ba9ad7
1 changed files with 50 additions and 140 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

190
db.py
View File

@ -2,6 +2,7 @@ import psycopg2
import config import config
import utils import utils
def connect(): def connect():
global cur, conn global cur, conn
try: try:
@ -12,9 +13,9 @@ def connect():
port=config.port) port=config.port)
except Exception as ex: except Exception as ex:
print(f"Error connecting to PostgreSQL: {ex}") print(f"Error connecting to PostgreSQL: {ex}")
cur = conn.cursor() cur = conn.cursor()
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute(""" cur.execute("""
CREATE TABLE IF NOT EXISTS image ( CREATE TABLE IF NOT EXISTS image (
@ -47,120 +48,77 @@ def connect():
client_ip INET client_ip INET
);""") );""")
conn.commit() conn.commit()
def get_cur(): def get_cur():
return conn.cursor() return conn.cursor()
def get_conn(): def get_conn():
return conn return conn
def get_one(sql, value): def get_one(sql, value):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(sql,(value,)) cur.execute(sql, (value,))
try: try:
return cur.fetchone()[0] return cur.fetchone()[0]
except: except:
return None return None
def add_conf_image(name, token): def add_conf_image(name, token):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
INSERT INTO image (image_name, token) INSERT INTO image (image_name, token)
VALUES (%s, %s) VALUES (%s, %s)
""",(name, token,)) """, (name, token,))
conn.commit() 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): def get_conf_image(token):
return get_one("SELECT image_name FROM image WHERE token = %s", 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): def get_conf_id(token):
return get_one("SELECT id FROM image WHERE token = %s", 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): def get_conf_id_name(name):
return get_one("SELECT id FROM image WHERE image_name = %s", 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): def add_user(username, password):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
INSERT INTO users (username, password) INSERT INTO users (username, password)
VALUES (%s, %s) VALUES (%s, %s)
""",(username, utils.hash_password(password),)) """, (username, utils.hash_password(password),))
conn.commit() conn.commit()
def get_user_pass(username, password): def get_user_pass(username, password):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
SELECT id FROM users WHERE username = %s AND password = %s SELECT id FROM users WHERE username = %s AND password = %s
""",(username, utils.hash_password(password),)) """, (username, utils.hash_password(password),))
try: try:
return cur.fetchone()[0] return cur.fetchone()[0]
except: except:
return None return None
def get_user_byid(id): def get_user_byid(id):
return get_one("SELECT id FROM users WHERE id = %s", 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): def get_user_bytoken(token):
return get_one("SELECT user_id FROM auth_tokens WHERE token = %s", 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): def add_auth_token(user_id):
token = utils.generate_auth_token() token = utils.generate_auth_token()
@ -169,18 +127,19 @@ def add_auth_token(user_id):
cur.execute(""" cur.execute("""
INSERT INTO auth_tokens (user_id, token, expires_on) INSERT INTO auth_tokens (user_id, token, expires_on)
VALUES (%s, %s, NOW() + INTERVAL '1 day') VALUES (%s, %s, NOW() + INTERVAL '1 day')
""",(user_id,token,)) """, (user_id, token,))
conn.commit() conn.commit()
return token return token
def login(username, password): def login(username, password):
user_id = get_user_pass(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:
return None return None
def get_image_allocation_all(): def get_image_allocation_all():
connect() connect()
with get_cur() as cur: with get_cur() as cur:
@ -202,85 +161,53 @@ def get_image_allocation(sql, value):
except: except:
return None return None
def get_image_allocation_time(token): def get_image_allocation_time(token):
image_id = get_conf_id(token) image_id = get_conf_id(token)
if image_id is None: if image_id is None:
return None return None
return get_one("SELECT last_access_time FROM image_allocation WHERE image_id = %s", image_id) 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): def get_image_allocation_time_id(id):
get_one("SELECT last_access_time FROM image_allocation WHERE id = %s", 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): 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
return get_one("SELECT last_access_time FROM image_allocation WHERE id = %s", id_image) 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): def get_image_allocation_clientip_id(id):
return get_one("SELECT client_ip FROM image_allocation WHERE id = %s", 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
def set_image_allocation(token, client_ip): def set_image_allocation(token, client_ip):
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() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
INSERT INTO image_allocation (image_id, client_ip, last_access_time) INSERT INTO image_allocation (image_id, client_ip, last_access_time)
VALUES (%s, %s, NOW()) VALUES (%s, %s, NOW())
""",(id_image,client_ip,)) """, (id_image, client_ip,))
conn.commit() conn.commit()
return token return token
def del_image_allocation_token(token): def del_image_allocation_token(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
return del_image_allocation_id_image(id_image) return del_image_allocation_id_image(id_image)
def del_image_allocation(sql, value): def del_image_allocation(sql, value):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
@ -291,40 +218,23 @@ def del_image_allocation(sql, value):
except: except:
return None return None
def del_image_allocation_id_image(image_id): def del_image_allocation_id_image(image_id):
return del_image_allocation("DELETE FROM image_allocation WHERE image_id = %s", 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) def del_image_allocation_id(id):
# connect() return del_image_allocation("DELETE FROM image_allocation WHERE id = %s", id)
# with get_cur() as cur:
# cur.execute("""
# DELETE FROM image_allocation WHERE id = %s def update_image_allocation_time(id):
# """,(id,))
# try:
# conn.commit()
# return True
# except:
# return None
def update_image_allocation_time(id):
connect() connect()
with get_cur() as cur: with get_cur() as cur:
cur.execute(""" cur.execute("""
UPDATE image_allocation SET last_access_time = NOW() WHERE id = %s UPDATE image_allocation SET last_access_time = NOW() WHERE id = %s
""",(id,)) """, (id,))
try: try:
conn.commit() conn.commit()
return True return True
except: except:
return None return None