transient.chat/models/database.py

125 lines
4.2 KiB
Python

import sqlite3
import os
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "transient_chat.db")
def get_db_connection():
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
return conn
def init_db():
with get_db_connection() as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS bans (
id INTEGER PRIMARY KEY AUTOINCREMENT,
type TEXT,
value TEXT UNIQUE
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS room_owners (
room TEXT PRIMARY KEY,
owner TEXT
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS user_ips (
username TEXT PRIMARY KEY,
ip TEXT
)
""")
conn.execute("""
CREATE TABLE IF NOT EXISTS public_rooms (
room TEXT PRIMARY KEY
)
""")
conn.commit()
def load_persisted_state(state_module):
"""Loads database state into the memory-based collections of state_module."""
init_db()
with get_db_connection() as conn:
# Load bans
for row in conn.execute("SELECT type, value FROM bans"):
if row["type"] == "ip":
state_module.banned_ips.add(row["value"])
elif row["type"] == "username":
state_module.banned_usernames.add(row["value"])
# Load room owners
for row in conn.execute("SELECT room, owner FROM room_owners"):
state_module.room_owners[row["room"]] = row["owner"]
# Load user IPs
for row in conn.execute("SELECT username, ip FROM user_ips"):
state_module.user_ips[row["username"]] = row["ip"]
# Load public rooms
for row in conn.execute("SELECT room FROM public_rooms"):
if row["room"] not in state_module.public_chat_rooms:
state_module.public_chat_rooms.append(row["room"])
def db_add_ban(ban_type, value):
try:
with get_db_connection() as conn:
conn.execute("INSERT OR IGNORE INTO bans (type, value) VALUES (?, ?)", (ban_type, value))
conn.commit()
except Exception as e:
print(f"DB Error db_add_ban: {e}")
def db_remove_ban(ban_type, value):
try:
with get_db_connection() as conn:
conn.execute("DELETE FROM bans WHERE type = ? AND value = ?", (ban_type, value))
conn.commit()
except Exception as e:
print(f"DB Error db_remove_ban: {e}")
def db_set_room_owner(room, owner):
try:
with get_db_connection() as conn:
conn.execute("INSERT OR REPLACE INTO room_owners (room, owner) VALUES (?, ?)", (room, owner))
conn.commit()
except Exception as e:
print(f"DB Error db_set_room_owner: {e}")
def db_delete_room_owner(room):
try:
with get_db_connection() as conn:
conn.execute("DELETE FROM room_owners WHERE room = ?", (room,))
conn.commit()
except Exception as e:
print(f"DB Error db_delete_room_owner: {e}")
def db_set_user_ip(username, ip):
try:
with get_db_connection() as conn:
conn.execute("INSERT OR REPLACE INTO user_ips (username, ip) VALUES (?, ?)", (username, ip))
conn.commit()
except Exception as e:
print(f"DB Error db_set_user_ip: {e}")
def db_delete_user_ip(username):
try:
with get_db_connection() as conn:
conn.execute("DELETE FROM user_ips WHERE username = ?", (username,))
conn.commit()
except Exception as e:
print(f"DB Error db_delete_user_ip: {e}")
def db_add_public_room(room):
try:
with get_db_connection() as conn:
conn.execute("INSERT OR IGNORE INTO public_rooms (room) VALUES (?)", (room,))
conn.commit()
except Exception as e:
print(f"DB Error db_add_public_room: {e}")
def db_remove_public_room(room):
try:
with get_db_connection() as conn:
conn.execute("DELETE FROM public_rooms WHERE room = ?", (room,))
conn.commit()
except Exception as e:
print(f"DB Error db_remove_public_room: {e}")