#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Coded by L1ghtM3n@protonmail.com

# Import modules
from base64 import b64encode, b64decode
# Import packages
from core.models.accounts import clients_db

# Get encryption key
def GetEncryptionKey(owner: str) -> str:
    response = clients_db.execute(
            sql="SELECT key FROM encryption WHERE owner = ?",
            values=[owner],
            commit_changes=False
        )
    # Skip if empty
    if len(response) == 0:
        return ""
    # Return values
    else:
        return response[0][0]

# Change encryption key
def SetEncryptionKey(owner: str, key: str):
    clients_db.execute(
            sql="UPDATE encryption SET key = ? WHERE owner = ?",
            values=[key, owner],
            commit_changes=True
        )

class RC4Cipher:
    """ RC4 encrypt/decrypt """
    def __init__(self, key):
        assert(isinstance(key, (bytes, bytearray)))

        # key scheduling
        S = list(range(0x100))
        j = 0
        for i in range(0x100):
            j = (S[i] + key[i % len(key)] + j) & 0xff
            S[i], S[j] = S[j], S[i]
        self.S = S

    def CryptBytes(self, data):
        """
        Encrypts/decrypts data (It's the same thing!)
        """
        assert(isinstance(data, (bytes, bytearray)))
        return bytes([a ^ b for a, b in zip(data, self._keystream_generator())])

    def Encrypt(self, string):
        """
        Encrypts string
        """
        return b64encode(self.CryptBytes(string.encode('utf8'))).decode('utf8')

    def Decrypt(self, string):
        """
        Decrypts string
        """
        return self.CryptBytes(b64decode(string)).decode('utf8')

    def _keystream_generator(self):
        """
        Generator that returns the bytes of keystream
        """
        S = self.S.copy()
        x = y = 0
        while True:
            x = (x + 1) & 0xff
            y = (S[x] + y) & 0xff
            S[x], S[y] = S[y], S[x]
            i = (S[x] + S[y]) & 0xff
            yield S[i]