latosa/auth.py

import json

from cryptography.fernet import Fernet

from db import DB


class Auth:

    def __init__(self):
        self.db = DB()

    def get_key(self, client_id):
        bind_params = {'client_id': client_id}
        return self.db.execute(
            'SELECT client_secret FROM auth WHERE client_id = :client_id',
            bind_params)[0]

    def set_key(self, client_id, client_secret):
        bind_params = {'client_id': client_id, 'client_secret': client_secret}
        self.db.execute(
            'INSERT INTO auth (client_id, client_secret) VALUES (:client_id, :client_secret)',
            bind_params)

    def decrypt(self, cyphertext, client_id) -> dict:
        return json.loads(
            Fernet(self.get_key(client_id)).decrypt(cyphertext).decode())

    def encrypt(self, session_key, data) -> bytes:
        return Fernet(session_key).encrypt(json.dumps(data).encode('utf-8'))