#!/usr/bin/env python3

from Crypto.Util.number import getPrime, bytes_to_long
from flask import Flask, request, jsonify, abort
from logging import getLogger
import secrets
import sqlite3
import os

E = 3
TOTAL_LEN = 252

app = Flask(__name__)
logger = getLogger(__name__)


def get_pub():
    # The public key seems to be retrieved from the boss' computer, which essentially
    # does this:
    while True:
        p = getPrime(1024)
        q = getPrime(1024)
        if p % E != 1 and q % E != 1:
            break
    
    N = p * q
    return N


def RSA(data):
    N = get_pub()
    logger.info(f"N = {N}")

    padding = secrets.token_bytes(TOTAL_LEN - len(data))
    logger.info(f"padding = '{padding.hex()}'") # Oops

    m = bytes_to_long(padding + data)
    c = pow(m, E, N)

    return str(c)


def encrypt_and_store(data):
    if not data:
        abort(400, description="Empty request body")
    if len(data) >= TOTAL_LEN:
        abort(400, description="Message too long")

    conn = sqlite3.connect(os.environ['DB_PATH'])
    with conn.cursor() as cur:
        cur.execute(
            "INSERT INTO new_messages (message, length) VALUES (?, ?)",
            (RSA(data), len(data)),
        )
        msg_id = cur.lastrowid
        conn.commit()

    return { "id": msg_id }


@app.route("/message", methods=["POST"])
def store_message():
    return jsonify(encrypt_and_store(request.data))


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000)