buranun/database.py

112 lines
3.1 KiB
Python

import enum
import random
import unicodedata
import sqlite3
from passlib.hash import argon2
import config
# ------------------------------------------------------------------
# General
# ------------------------------------------------------------------
class userstatus(enum.Enum):
# These will be stored in the database, be mindful of not changing the numbers
deleted = 0
normal = 1
admin = 2
csprng = random.SystemRandom()
def connect():
"""Connect to the database
Requires config.load() to have been called beforehand"""
return sqlite3.connect(config.database_file)
# ------------------------------------------------------------------
# Users
# ------------------------------------------------------------------
def add_user(db, *, username, password, email, parent, status):
"""Add a user to the database
Will not commit the changes itself, so run .commit() on the database object yourself"""
global csprgn
assert type(username) == str
assert type(password) == str
assert type(email) == str
assert type(parent) == int or parent is None
assert status in userstatus
# Generate a user ID. SQLite uses 64 bit signed ints, so generate at max 2⁶³-1
userid = csprng.randrange(2**63)
# Unicode normalize the username
username = unicodedata.normalize('NFKC', username)
# First unicode normalize the password, then hash it with argon2
password = unicodedata.normalize('NFKC', password)
password = argon2.hash(password)
# Convert status into an int for storage
status = status.value
# Add the user into the database
cursor = db.cursor()
cursor.execute('PRAGMA foreign_keys = ON;') # Fail if we insert a user with bogus parent field
cursor.execute('INSERT INTO users VALUES (?, ?, ?, ?, ?, ?, ?);', (userid, parent, status, password, username, email, ''))
def initialize_users(db, admin_user, admin_password):
"""Creates a bare-bones user table with only admin user
This should never be run outside of the initialization script"""
cursor = db.cursor()
cursor.execute('''CREATE TABLE users (
id integer NOT NULL PRIMARY KEY,
parent integer,
status integer NOT NULL,
password text NOT NULL,
username text NOT NULL,
email text NOT NULL,
comment text NOT NULL,
FOREIGN KEY(parent) REFERENCES users(id)
);''')
add_user(db, username = admin_user, password = admin_password, email = '', parent = None, status = userstatus.admin)
db.commit()
# ------------------------------------------------------------------
# Boards
# ------------------------------------------------------------------
def list_boards(db):
# TODO: Implement this
...
def initialize_boards(db, boards):
"""Creates a table of boards
This should never be run outside of the initialization script"""
cursor = db.cursor()
cursor.execute('''CREATE TABLE boards (
id integer NOT NULL PRIMARY KEY,
name text NOT NULL
);''')
# .executemany() wants them in the format [("board1",), ("board2",), …]
boards = [(board_name,) for board_name in boards]
# Use NULL to have SQLite generate the IDs automatically
cursor.executemany('INSERT INTO boards VALUES (NULL, ?);', boards)
db.commit()