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

341 lines
9.2 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
import machines
2023-04-19 12:10:31 +00:00
import images
2023-04-18 12:51:04 +00:00
2023-04-21 08:40:09 +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:
2023-04-18 12:11:48 +00:00
print(f"Error connecting to PostgreSQL: {ex}")
2023-04-18 12:51:04 +00:00
2023-04-07 11:25:24 +00:00
cur = conn.cursor()
2023-04-18 12:51:04 +00:00
2023-04-07 11:25:24 +00:00
with conn.cursor() as cur:
2023-04-21 08:59:34 +00:00
cur.execute("SET TIMEZONE = 'UTC';")
2023-04-19 09:54:08 +00:00
conn.commit()
2023-04-07 11:25:24 +00:00
cur.execute("""
2023-04-14 08:41:07 +00:00
CREATE TABLE IF NOT EXISTS image (
2023-04-07 11:25:24 +00:00
id SERIAL PRIMARY KEY,
image_name VARCHAR(255) NOT NULL,
token VARCHAR(255) NOT NULL,
2023-04-20 14:29:01 +00:00
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
2023-04-21 08:40:09 +00:00
password VARCHAR(128) NOT NULL,
2023-04-20 14:29:01 +00:00
vpn_ip INET
2023-04-13 09:52:51 +00:00
);""")
2023-04-19 09:54:08 +00:00
conn.commit()
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
);""")
2023-04-19 09:54:08 +00:00
conn.commit()
2023-04-13 09:52:51 +00:00
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-19 09:54:08 +00:00
conn.commit()
2023-04-14 08:41:07 +00:00
cur.execute("""
CREATE TABLE IF NOT EXISTS image_allocation (
id SERIAL PRIMARY KEY,
image_id INTEGER NOT NULL REFERENCES image(id),
2023-04-19 09:54:08 +00:00
allocation_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
2023-04-14 08:41:07 +00:00
last_access_time TIMESTAMP,
2023-04-19 10:01:53 +00:00
client_ip_local INET,
client_ip_vpn INET
2023-04-14 08:41:07 +00:00
);""")
2023-04-07 11:25:24 +00:00
conn.commit()
2023-04-18 12:51:04 +00:00
2023-04-13 09:52:51 +00:00
def get_cur():
2023-04-07 11:25:24 +00:00
return conn.cursor()
2023-04-18 12:51:04 +00:00
2023-04-13 09:52:51 +00:00
def get_conn():
2023-04-07 11:25:24 +00:00
return conn
2023-04-18 12:51:04 +00:00
2023-04-18 12:30:56 +00:00
def get_one(sql, value):
connect()
with get_cur() as cur:
2023-04-18 12:51:04 +00:00
cur.execute(sql, (value,))
2023-04-18 12:30:56 +00:00
try:
return cur.fetchone()[0]
except:
return None
2023-04-18 12:51:04 +00:00
2023-04-21 08:40:09 +00:00
def add_conf_image(name, token, ip, password):
2023-04-13 09:52:51 +00:00
connect()
with get_cur() as cur:
2023-04-07 11:25:24 +00:00
cur.execute("""
2023-04-21 08:40:09 +00:00
INSERT INTO image (image_name, token, vpn_ip, password)
VALUES (%s, %s, %s, %s)
""", (name, token, ip, password, ))
2023-04-07 11:25:24 +00:00
conn.commit()
2023-04-18 12:51:04 +00:00
2023-04-18 12:11:48 +00:00
def get_conf_image(token):
2023-04-18 12:30:56 +00:00
return get_one("SELECT image_name FROM image WHERE token = %s", token)
2023-04-18 12:51:04 +00:00
2023-04-21 08:40:09 +00:00
def get_conf_password(token):
return get_one("SELECT password FROM image WHERE token = %s", token)
2023-04-19 13:31:25 +00:00
def get_conf_image_id(id):
return get_one("SELECT image_name FROM image WHERE id = %s", id)
2023-04-21 08:40:09 +00:00
2023-04-14 08:41:07 +00:00
def get_conf_id(token):
2023-04-18 12:30:56 +00:00
return get_one("SELECT id FROM image WHERE token = %s", token)
2023-04-18 12:51:04 +00:00
2023-04-14 12:02:01 +00:00
def get_conf_id_name(name):
2023-04-18 12:30:56 +00:00
return get_one("SELECT id FROM image WHERE image_name = %s", name)
2023-04-18 12:51:04 +00:00
2023-04-13 09:52:51 +00:00
def add_user(username, password):
connect()
with get_cur() as cur:
cur.execute("""
INSERT INTO users (username, password)
VALUES (%s, %s)
2023-04-18 12:51:04 +00:00
""", (username, utils.hash_password(password),))
2023-04-13 09:52:51 +00:00
conn.commit()
2023-04-18 12:51:04 +00:00
2023-04-18 12:11:48 +00:00
def get_user_pass(username, password):
2023-04-13 09:52:51 +00:00
connect()
with get_cur() as cur:
cur.execute("""
SELECT id FROM users WHERE username = %s AND password = %s
2023-04-18 12:51:04 +00:00
""", (username, utils.hash_password(password),))
2023-04-13 09:52:51 +00:00
try:
return cur.fetchone()[0]
except:
return None
2023-04-18 12:11:48 +00:00
2023-04-18 12:51:04 +00:00
2023-04-13 09:52:51 +00:00
def get_user_byid(id):
2023-04-18 12:30:56 +00:00
return get_one("SELECT id FROM users WHERE id = %s", id)
2023-04-18 12:51:04 +00:00
2023-04-13 09:52:51 +00:00
def get_user_bytoken(token):
2023-04-18 12:30:56 +00:00
return get_one("SELECT user_id FROM auth_tokens WHERE token = %s", token)
2023-04-13 09:52:51 +00:00
def add_auth_token(user_id):
token = utils.generate_auth_token()
connect()
with get_cur() as cur:
cur.execute("""
2023-04-13 10:31:13 +00:00
INSERT INTO auth_tokens (user_id, token, expires_on)
2023-04-19 09:54:08 +00:00
VALUES (%s, %s, CURRENT_TIMESTAMP + INTERVAL '1 day')
2023-04-18 12:51:04 +00:00
""", (user_id, token,))
2023-04-13 09:52:51 +00:00
conn.commit()
return token
2023-04-18 12:51:04 +00:00
2023-04-19 10:27:09 +00:00
def del_auth_token(token):
connect()
with get_cur() as cur:
cur.execute("DELETE FROM auth_tokens WHERE token = %s", (token, ))
try:
conn.commit()
return True
except:
return None
2023-04-13 09:52:51 +00:00
def login(username, password):
2023-04-18 12:11:48 +00:00
user_id = get_user_pass(username, password)
2023-04-13 09:52:51 +00:00
if user_id is not None:
return add_auth_token(user_id)
else:
return None
2023-04-18 12:51:04 +00:00
2023-04-21 08:40:09 +00:00
def get_machines():
connect()
with get_cur() as cur:
cur.execute("""
2023-04-24 06:27:35 +00:00
SELECT image_id, allocation_time, client_ip_vpn,
client_ip_local FROM image_allocation""")
try:
machinesall = machines.MachineManager()
for row in cur.fetchall():
2023-04-21 08:40:09 +00:00
token = get_one(
"SELECT token FROM image WHERE id = %s", row[0])
image_name = get_one(
"SELECT image_name FROM image WHERE id = %s", row[0])
machine = machines.Machine(
2023-04-24 06:27:35 +00:00
token, image_name, start_time=row[1], ipvpn=row[2],
iplocal=row[3], username="root", password="")
machinesall.add_machine(machine)
return machinesall
except:
return None
2023-04-21 08:40:09 +00:00
2023-04-19 12:10:31 +00:00
def get_images():
connect()
with get_cur() as cur:
cur.execute("""
2023-04-20 14:29:01 +00:00
SELECT id, token, image_name, vpn_ip FROM image""")
2023-04-19 12:10:31 +00:00
try:
images_all = images.ImageManager()
for row in cur.fetchall():
2023-04-21 08:40:09 +00:00
image = images.Image(
id=row[0], token=row[1], name=row[2], vpn_ip=row[3])
2023-04-19 12:10:31 +00:00
images_all.add_image(image)
return images_all
except:
return None
2023-04-21 08:40:09 +00:00
2023-04-19 13:31:25 +00:00
def del_image(image_id):
connect()
with get_cur() as cur:
cur.execute("DELETE FROM image WHERE id = %s", (image_id,))
try:
conn.commit()
return True
except:
return None
2023-04-18 12:51:04 +00:00
2023-04-21 08:40:09 +00:00
def get_image_allocation_all_id():
2023-04-14 08:41:07 +00:00
connect()
with get_cur() as cur:
cur.execute("""
2023-04-14 09:54:28 +00:00
SELECT id FROM image_allocation""")
2023-04-14 08:41:07 +00:00
try:
results = [list(row) for row in cur.fetchall()]
return results
except:
return None
2023-04-18 12:11:48 +00:00
def get_image_allocation_all():
connect()
with get_cur() as cur:
cur.execute("""
SELECT * FROM image_allocation""")
try:
results = [list(row) for row in cur.fetchall()]
return results
except:
return None
2023-04-21 08:40:09 +00:00
def get_image_allocation(image_id):
return get_one("SELECT id FROM image_allocation WHERE image_id = %s", image_id)
2023-04-14 08:41:07 +00:00
2023-04-18 12:51:04 +00:00
2023-04-18 12:11:48 +00:00
def get_image_allocation_time(token):
image_id = get_conf_id(token)
if image_id is None:
return None
2023-04-18 12:30:56 +00:00
return get_one("SELECT last_access_time FROM image_allocation WHERE image_id = %s", image_id)
2023-04-18 12:11:48 +00:00
2023-04-18 12:51:04 +00:00
2023-04-14 09:24:32 +00:00
def get_image_allocation_time_id(id):
return get_one("SELECT last_access_time FROM image_allocation WHERE id = %s", id)
2023-04-18 12:51:04 +00:00
2023-04-14 08:41:07 +00:00
def get_image_allocation_clientip(token):
id_image = get_conf_id(token)
if id_image is None:
return None
2023-04-18 12:51:04 +00:00
2023-04-18 12:30:56 +00:00
return get_one("SELECT last_access_time FROM image_allocation WHERE id = %s", id_image)
2023-04-18 12:51:04 +00:00
2023-04-14 08:41:07 +00:00
2023-04-19 10:01:53 +00:00
def get_image_allocation_clientip_id_vpn(id):
return get_one("SELECT client_ip_vpn FROM image_allocation WHERE id = %s", id)
2023-04-18 12:51:04 +00:00
2023-04-14 08:41:07 +00:00
def set_image_allocation(token, client_ip):
id_image = get_conf_id(token)
if id_image is None:
return None
2023-04-18 12:51:04 +00:00
2023-04-14 08:41:07 +00:00
connect()
with get_cur() as cur:
cur.execute("""
2023-04-19 10:01:53 +00:00
INSERT INTO image_allocation (image_id, client_ip_local, last_access_time)
2023-04-19 09:54:08 +00:00
VALUES (%s, %s, CURRENT_TIMESTAMP)
2023-04-18 12:51:04 +00:00
""", (id_image, client_ip,))
2023-04-14 08:41:07 +00:00
conn.commit()
2023-04-14 09:24:32 +00:00
return token
2023-04-18 12:51:04 +00:00
2023-04-14 09:24:32 +00:00
def del_image_allocation_token(token):
id_image = get_conf_id(token)
if id_image is None:
2023-04-18 12:51:04 +00:00
return None
2023-04-14 12:22:08 +00:00
return del_image_allocation_id_image(id_image)
2023-04-14 09:24:32 +00:00
2023-04-18 12:51:04 +00:00
2023-04-18 12:21:34 +00:00
def del_image_allocation(sql, value):
2023-04-14 12:19:37 +00:00
connect()
with get_cur() as cur:
2023-04-18 12:21:34 +00:00
cur.execute(sql, (value, ))
2023-04-14 12:19:37 +00:00
try:
conn.commit()
return True
except:
return None
2023-04-18 12:51:04 +00:00
2023-04-18 12:11:48 +00:00
def del_image_allocation_id_image(image_id):
2023-04-19 08:22:38 +00:00
return del_image_allocation("DELETE FROM image_allocation WHERE image_id = %s", image_id)
2023-04-18 12:51:04 +00:00
def del_image_allocation_id(id):
return del_image_allocation("DELETE FROM image_allocation WHERE id = %s", id)
def update_image_allocation_time(id):
2023-04-14 09:24:32 +00:00
connect()
with get_cur() as cur:
cur.execute("""
2023-04-19 09:54:08 +00:00
UPDATE image_allocation SET last_access_time = CURRENT_TIMESTAMP WHERE id = %s
2023-04-18 12:51:04 +00:00
""", (id,))
2023-04-19 08:14:39 +00:00
try:
conn.commit()
return True
except:
return None
2023-04-21 08:40:09 +00:00
2023-04-19 10:01:53 +00:00
def update_image_allocation_ip_vpn(token, ip):
2023-04-19 08:22:38 +00:00
image_id = get_conf_id(token)
if image_id is None:
return None
2023-04-19 08:14:39 +00:00
connect()
with get_cur() as cur:
cur.execute("""
2023-04-19 10:01:53 +00:00
UPDATE image_allocation SET client_ip_vpn = %s WHERE image_id = %s
2023-04-19 08:22:38 +00:00
""", (ip, image_id,))
2023-04-14 09:24:32 +00:00
try:
conn.commit()
return True
except:
2023-04-21 08:40:09 +00:00
return None