129 lines
3.6 KiB
Python
129 lines
3.6 KiB
Python
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 vpn (
|
|
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
|
|
);""")
|
|
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 vpn (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 vpn 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
|
|
|