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

129 lines
3.5 KiB
Python
Raw Normal View History

2023-04-07 11:25:24 +00:00
import psycopg2
import config
2023-04-13 09:52:51 +00:00
import utils
2023-04-07 11:25:24 +00:00
2023-04-13 09:52:51 +00:00
def connect():
2023-04-07 11:25:24 +00:00
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 vpn (
id SERIAL PRIMARY KEY,
image_name VARCHAR(255) NOT NULL,
token VARCHAR(255) NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
2023-04-13 09:52:51 +00:00
);""")
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("""
2023-04-13 09:57:50 +00:00
CREATE TABLE IF NOT EXISTS auth_tokens (
2023-04-13 09:52:51 +00:00
id SERIAL PRIMARY KEY,
user_id INTEGER NOT NULL REFERENCES users(id),
token VARCHAR(64) NOT NULL,
2023-04-13 09:55:02 +00:00
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
2023-04-13 09:52:51 +00:00
expires_on TIMESTAMP NOT NULL
);""")
2023-04-07 11:25:24 +00:00
conn.commit()
2023-04-13 09:52:51 +00:00
def get_cur():
2023-04-07 11:25:24 +00:00
return conn.cursor()
2023-04-13 09:52:51 +00:00
def get_conn():
2023-04-07 11:25:24 +00:00
return conn
2023-04-13 09:52:51 +00:00
def add_conf_image(name, token):
connect()
with get_cur() as cur:
2023-04-07 11:25:24 +00:00
cur.execute("""
INSERT INTO vpn (image_name, token)
VALUES (%s, %s)
""",(name, token,))
conn.commit()
2023-04-13 09:52:51 +00:00
def get_conf_image(token):
connect()
with get_cur() as cur:
2023-04-07 11:32:49 +00:00
cur.execute("""
2023-04-07 11:25:24 +00:00
SELECT image_name FROM vpn WHERE token = %s
2023-04-07 11:32:49 +00:00
""",(token,))
2023-04-13 09:52:51 +00:00
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
2023-04-07 11:25:24 +00:00
2023-04-13 09:52:51 +00:00
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)
VALUES (%s, %s)
""",(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