gistfile1.txt

class KeyManager:
    def __init__(self, db_name="demo_polynomial.db"):
        self.db_name = db_name
        self._create_tables()
        self.crypto_string = None
        self.crypto_csv = None
        self._initialize_demo_data()

    def _get_connection(self):
        return sqlite3.connect(self.db_name)

    def _create_tables(self):
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS UserIds (
                user_id INTEGER PRIMARY KEY AUTOINCREMENT,
                api_key TEXT UNIQUE,
                assigned_poly_key TEXT
            );
            """)
            cursor.execute("""
            CREATE TABLE IF NOT EXISTS decryptKeys (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                key_value TEXT UNIQUE
            );
            """)
            conn.commit()

    def _initialize_demo_data(self):
        string_demo = "Mensagem secreta mestre transmitida com seguranca!"
        self.crypto_string = PolynomialMultiKeyEncryption(string_demo.encode('utf-8'))

        csv_demo = "id,nome,cargo\n1,Alice,Engenheira\n2,Bruno,Cientista de Dados"
        self.crypto_csv = PolynomialMultiKeyEncryption(csv_demo.encode('utf-8'))

    def generate_and_register_keys(self, amount: int):
        new_keys = []
        for _ in range(amount):
            poly_key = self.crypto_string.generate_new_key()
            new_keys.append((poly_key,))
            
        with self._get_connection() as conn:
            cursor = conn.cursor()
            try:
                cursor.executemany("INSERT INTO decryptKeys (key_value) VALUES (?);", new_keys)
                conn.commit()
                return True
            except sqlite3.Error:
                conn.rollback()
                return False

    def get_or_assign_key(self, api_key: str):
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT assigned_poly_key FROM UserIds WHERE api_key = ?;", (api_key,))
            existing = cursor.fetchone()
            if existing and existing[0]:
                return existing[0]

            cursor.execute("SELECT id, key_value FROM decryptKeys ORDER BY id ASC LIMIT 1;")
            result = cursor.fetchone()
            if not result:
                return None
            
            key_id, poly_key_value = result
            cursor.execute("INSERT OR REPLACE INTO UserIds (api_key, assigned_poly_key) VALUES (?, ?);", (api_key, poly_key_value))
            cursor.execute("DELETE FROM decryptKeys WHERE id = ?;", (key_id,))
            conn.commit()
            return poly_key_value

    def check_key_valid(self, api_key: str, key_to_check: str) -> bool:
        with self._get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT assigned_poly_key FROM UserIds WHERE api_key = ?;", (api_key,))
            result = cursor.fetchone()
            return result and result[0] == key_to_check


manager = KeyManager()