Zdalne-systemy-z-kit-uVPN/db.py

import psycopg2
import config
import utils

def connect():
    global cur, conn
    try:
        conn = psycopg2.connect(database=config.database,
                                host=config.host,
                                user=config.user,
                                password=config.password,
                                port=config.port)
    except Exception as ex:
        print(f"Error connecting to PostgreSQL: {e}")
        
    cur = conn.cursor()
    
    with conn.cursor() as cur:
        cur.execute("""
            CREATE TABLE IF NOT EXISTS image (
                id SERIAL PRIMARY KEY,
                image_name VARCHAR(255) NOT NULL,
                token VARCHAR(255) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );""")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id SERIAL PRIMARY KEY,
                username VARCHAR(50) UNIQUE NOT NULL,
                password VARCHAR(256) NOT NULL,
                created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
            );""")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS auth_tokens (
                id SERIAL PRIMARY KEY,
                user_id INTEGER NOT NULL REFERENCES users(id),
                token VARCHAR(64) NOT NULL,
                created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
                expires_on TIMESTAMP NOT NULL
            );""")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS image_allocation (
                id SERIAL PRIMARY KEY,
                image_id INTEGER NOT NULL REFERENCES image(id),
                allocation_time TIMESTAMP NOT NULL DEFAULT NOW(),
                last_access_time TIMESTAMP,
                client_ip INET
            );""")
        conn.commit()
    
def get_cur():
    return conn.cursor()

def get_conn():
    return conn

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,))
        conn.commit()
        
def get_conf_image(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):
    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 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),))
        conn.commit()
        
def get_user(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),))
        try:
            return cur.fetchone()[0]
        except:
            return None
        
def get_user_byid(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):
    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()
    connect()
    with get_cur() as cur:
        cur.execute("""
            INSERT INTO auth_tokens (user_id, token, expires_on)
            VALUES (%s, %s, NOW() + INTERVAL '1 day')
        """,(user_id,token,))
        conn.commit()
    return token
    
def login(username, password):
    user_id = get_user(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:
        cur.execute("""
            SELECT id FROM image_allocation""")
        try:
            results = [list(row) for row in cur.fetchall()]
            return results
        except:
            return None

def get_image_allocation_time(token):
    id_image = get_conf_id(token)
    if id_image is None:
        return None
    
    connect()
    with get_cur() as cur:
        cur.execute("""
            SELECT last_access_time FROM image_allocation WHERE image_id = %s 
        """,(id_image,))
        try:
            return cur.fetchone()[0]
        except:
            return None

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
        
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
        
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

def get_image_allocation_clientip_id(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):
    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,))
        conn.commit()
    return token

def del_image_allocation_token(token):
    id_image = get_conf_id(token)
    if id_image is None:
        return None 
    
    connect()
    with get_cur() as cur:
        cur.execute("""
            DELETE FROM image_allocation WHERE id_image = %s 
        """,(id_image,))
        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 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,))
        try:
            conn.commit()
            return True
        except:
            return None