database.py

import base64
import dataclasses
import json
import logging

import asyncpg

from cryptography.exceptions import InvalidSignature
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

import config
from structs import Spoiler
from database_id import UUID


pool: asyncpg.pool.Pool
logger = logging.getLogger('db')


def v1_derive_key(uuid, salt):
  """derives a key from a uuid+unique salt using scrypt"""
  return base64.urlsafe_b64encode(
    Scrypt(
      salt=salt,
      length=32,
      n=2**10,
      r=8,
      p=1,
      backend=default_backend()
    ).derive(uuid.encode('ascii'))
  )


def v1_hash_uuid(uuid):
  """
  hashes a uuid using SHA256 (these are the primary keys of the database)
  we can't use a unique salt here because we need the hash to find the row
  """
  digest = hashes.Hash(hashes.SHA256(), backend=default_backend())
  digest.update((uuid + config.HASH_PEPPER).encode())
  return digest.finalize()


def v2_hash_uuid(uuid):
  """
  Converts a (variable length) uuid into a database key and encryption key by
  splitting the SHA512 hash
  """
  digest = hashes.Hash(hashes.SHA512(), backend=default_backend())
  digest.update((uuid + config.HASH_PEPPER).encode())
  digest = digest.finalize()
  return digest[:32], base64.urlsafe_b64encode(digest[32:])


async def init():
  global pool
  logger.info('Creating connection pool...')
  pool = await asyncpg.create_pool(
    database=config.DB_NAME,
    user=config.DB_USERNAME,
    host=config.DB_HOST,
    password=config.DB_PASSWORD
  )

  logger.info('Creating tables...')
  async with pool.acquire() as con:
    await con.execute('''
      CREATE TABLE IF NOT EXISTS spoilers (
        hash BYTEA PRIMARY KEY,
        timestamp INTEGER DEFAULT date_part('epoch', now()),
        salt BYTEA,
        token BYTEA
      )
    ''')
    await con.execute('''
      CREATE TABLE IF NOT EXISTS spoilers_v2 (
        hash BYTEA PRIMARY KEY,
        timestamp INTEGER DEFAULT date_part('epoch', now()),
        token BYTEA,
        owner BIGINT
      )
    ''')
  await pool.execute('VACUUM;')
  await pool.execute('TRUNCATE spoilers_v2;')
  logger.info('Initialized.')


async def _get_spoiler_v1(uuid: str):
  """
  Gets a spoiler from the v1 schema
  """
  # try to find uuid by hash in the database
  db_hash = v1_hash_uuid(uuid)
  data = await pool.fetchrow(
    'SELECT salt, token FROM spoilers WHERE hash=$1',
    db_hash
  )

  if not data:
    return None

  # Decrypt the data and decode it
  try:
    key = v1_derive_key(uuid, data['salt'])
    data = Fernet(key).decrypt(data['token'])
  except InvalidSignature:
    # this shouldn't happen unless someone messes with the database
    return None

  return Spoiler(**json.loads(data))


async def get_spoiler(uuid: UUID):
  db_hash, key = v2_hash_uuid(uuid.db_key)

  data = await pool.fetchval(
    'SELECT token FROM spoilers_v2 WHERE hash=$1',
    db_hash
  )

  if not data:
    return await _get_spoiler_v1(uuid.db_key)

  # Decrypt the data and decode it
  try:
    data = Fernet(key).decrypt(data)
  except InvalidSignature:
    # this shouldn't happen unless someone messes with the database
    return None
  return Spoiler(**json.loads(data))


async def insert_spoiler(uuid: UUID, spoiler: Spoiler, owner_id: int):
  data = json.dumps(dataclasses.asdict(spoiler)).encode()

  db_hash, key = v2_hash_uuid(uuid.db_key)
  token = Fernet(key).encrypt(data)

  await pool.execute(
    'INSERT INTO spoilers_v2 (hash, token, owner) VALUES ($1, $2, $3)',
    db_hash, token, owner_id
  )