from transforms.base import *

def rol(n, rot, width):
    return (
        (n << rot % width) & (2**width-1) |
        ((n & (2**width-1)) >> (width-(rot % width)))
    )

def ror(n, rot, width):
    return (
        ((n & (2**width-1)) >> rot % width) |
        (n << (width - (rot % width)) & (2**width-1))
    )

class AesDecrypt(Transform):
    """
    AES cipher decryption

    """
    category = "crypto (block)"
    name = "aes decrypt"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb", "ctr", "gcm")="ecb", key:bytes=b"aaaaaaaaaaaaaaaa", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", unpad:bool=False):
        from Cryptodome.Cipher import AES
        from Cryptodome.Util import Padding, Counter
        if mode == "ecb":
            decryptor = AES.new(key, AES.MODE_ECB)
        elif mode == "cbc":
            decryptor = AES.new(key, AES.MODE_CBC, iv=iv)
        elif mode == "cfb":
            decryptor = AES.new(key, AES.MODE_CFB, iv=iv)
        elif mode == "ctr":
            decryptor = AES.new(key, AES.MODE_CTR, counter=Counter.new(128, initial_value=int.from_bytes(iv, byteorder='big')))
        elif mode == "gcm":
            decryptor = AES.new(key, AES.MODE_GCM, nonce=iv)
        data = decryptor.decrypt(data)
        if unpad:
            data = Padding.unpad(data, 16)
        return data


class AesDecryptPBKDF2(Transform):
    """
    AES cipher decryption, but password and IV are generated using PBKDF2/Rfc2898
    """
    category = "crypto (block)"
    name = "aes decrypt PBKDF2"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, password:bytes=b"aaaaaaaaaaaaaaaa", salt:bytes=b"\x01\x02\x03\x04\x05\x06\x07\x08", rounds:int=10000, hash_password:["no", "md5", "sha1", "sha256", "sha512"]="no", mode:("ecb","cbc","cfb", "ctr", "gcm")="ecb", unpad:bool=False):
        from Cryptodome.Protocol.KDF import PBKDF2
        import hashlib
        if hash_password == "md5":
            password = hashlib.md5(password).digest()
        elif hash_password == "sha1":
            password = hashlib.sha1(password).digest()
        elif hash_password == "sha256":
            password = hashlib.sha512(password).digest()            
        elif hash_password == "sha512":
            password = hashlib.sha256(password).digest()
        keyiv = PBKDF2(password, salt, 32 + 16, count=rounds)
        return AesDecrypt().run(data, mode, keyiv[:32], keyiv[32:])


class AesEncrypt(Transform):
    """
    AES cipher encryption
    """
    category = "crypto (block)"
    name = "aes encrypt"
    icon = "wxART_CRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb", "ctr", "gcm")="ecb", key:bytes=b"aaaaaaaaaaaaaaaa", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", pad:bool=False):
        from Cryptodome.Cipher import AES
        from Cryptodome.Util import Padding, Counter
        if pad:
            data = Padding.pad(data, 16)
        if mode == "ecb":
            cryptor = AES.new(key, AES.MODE_ECB)
        elif mode == "cbc":
            cryptor = AES.new(key, AES.MODE_CBC, iv=iv)
        elif mode == "cfb":
            cryptor = AES.new(key, AES.MODE_CFB, iv=iv)
        elif mode == "ctr":
            cryptor = AES.new(key, AES.MODE_CTR, counter=Counter.new(128, initial_value=int.from_bytes(iv, byteorder='big')))
        elif mode == "gcm":
            cryptor = AES.new(key, AES.MODE_GCM, nonce=iv)
        return cryptor.encrypt(data)


class BlowfishDecrypt(Transform):
    """
    Blowfish cipher decryption
    """
    category = "crypto (block)"
    name = "blowfish decrypt"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"\x00\x00\x00\x00", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", unpad:bool=False):
        from Cryptodome.Cipher import Blowfish
        from Cryptodome.Util import Padding
        if mode == "ecb":
            decryptor = Blowfish.new(key, Blowfish.MODE_ECB)
        elif mode == "cbc":
            decryptor = Blowfish.new(key, Blowfish.MODE_CBC, iv=iv)
        elif mode == "cfb":
            decryptor = Blowfish.new(key, Blowfish.MODE_CFB, iv=iv)
        data = decryptor.decrypt(data)
        if unpad:
            data = Padding.unpad(data, 8)
        return data    


class BlowfishEncrypt(Transform):
    """
    Blowfish cipher encryption
    """
    category = "crypto (block)"
    name = "blowfish encrypt"
    icon = "wxART_CRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"\x00\x00\x00\x00", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", pad:bool=False):
        from Cryptodome.Cipher import Blowfish
        from Cryptodome.Util import Padding
        if pad:
            data = Padding.pad(data, 8)
        if mode == "ecb":
            cryptor = Blowfish.new(key, Blowfish.MODE_ECB)
        elif mode == "cbc":
            cryptor = Blowfish.new(key, Blowfish.MODE_CBC, iv=iv)
        elif mode == "cfb":
            cryptor = Blowfish.new(key, Blowfish.MODE_CFB, iv=iv)
        return cryptor.encrypt(data)    


class Rc2Decrypt(Transform):
    """
    RC2 cipher decryption
    """
    category = "crypto (block)"
    name = "rc2 decrypt"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"aaaaaaaa", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", effective_keylen:int=128, unpad:bool=False):
        from Cryptodome.Cipher import ARC2
        from Cryptodome.Util import Padding
        options = {}
        if effective_keylen:
            options["effective_keylen"] = effective_keylen
        if mode == "ecb":
            decryptor = ARC2.new(key, ARC2.MODE_ECB, **options)
        elif mode == "cbc":
            decryptor = ARC2.new(key, ARC2.MODE_CBC, iv=iv, **options)
        elif mode == "cfb":
            decryptor = ARC2.new(key, ARC2.MODE_CFB, iv=iv, **options)

        data = decryptor.decrypt(data)
        if unpad:
            data = Padding.unpad(data, 8)
        return data    


class Rc2Encrypt(Transform):
    """
    RC2 cipher encryption
    """
    category = "crypto (block)"
    name = "rc2 encrypt"
    icon = "wxART_CRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"aaaaaaaa", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", effective_keylen:int=128, pad:bool=False):
        from Cryptodome.Cipher import ARC2
        from Cryptodome.Util import Padding
        if pad:
            data = Padding.pad(data, 8)
        options = {}
        if effective_keylen:
            options["effective_keylen"] = effective_keylen            
        if mode == "ecb":
            cryptor = ARC2.new(key, ARC2.MODE_ECB, **options)
        elif mode == "cbc":
            cryptor = ARC2.new(key, ARC2.MODE_CBC, iv=iv, **options)
        elif mode == "cfb":
            cryptor = ARC2.new(key, ARC2.MODE_CFB, iv=iv, **options)
        return cryptor.encrypt(data)        


class DesDecrypt(Transform):
    """
    DES cipher decryption
    """
    category = "crypto (block)"
    name = "des decrypt"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"aaaaaaaa", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", unpad:bool=False):
        from Cryptodome.Cipher import DES
        from Cryptodome.Util import Padding
        if mode == "ecb":
            decryptor = DES.new(key, DES.MODE_ECB)
        elif mode == "cbc":
            decryptor = DES.new(key, DES.MODE_CBC, iv=iv)
        elif mode == "cfb":
            decryptor = DES.new(key, DES.MODE_CFB, iv=iv)
        data = decryptor.decrypt(data)
        if unpad:
            data = Padding.unpad(data, 8)
        return data    


class DesEncrypt(Transform):
    """
    DES cipher encryption
    """
    category = "crypto (block)"
    name = "des encrypt"
    icon = "wxART_CRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"aaaaaaaa", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", pad:bool=False):
        from Cryptodome.Cipher import DES
        from Cryptodome.Util import Padding
        if pad:
            data = Padding.pad(data, 8)
        if mode == "ecb":
            cryptor = DES.new(key, DES.MODE_ECB)
        elif mode == "cbc":
            cryptor = DES.new(key, DES.MODE_CBC, iv=iv)
        elif mode == "cfb":
            cryptor = DES.new(key, DES.MODE_CFB, iv=iv)
        return cryptor.encrypt(data)        

class Des3Decrypt(Transform):
    """
    3DES cipher decryption
    key must be 16 or 24 bytes
    """
    category = "crypto (block)"
    name = "3des decrypt"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"0123456789abcdef", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", unpad:bool=False):
        from Cryptodome.Cipher import DES3, DES
        from Cryptodome.Util import Padding
        if mode == "ecb":
            decryptor = DES3.new(key, DES.MODE_ECB)
        elif mode == "cbc":
            decryptor = DES3.new(key, DES.MODE_CBC, iv=iv)
        elif mode == "cfb":
            decryptor = DES3.new(key, DES.MODE_CFB, iv=iv)
        data = decryptor.decrypt(data)
        if unpad:
            data = Padding.unpad(data, 8)
        return data    


class Des3Encrypt(Transform):
    """
    3DES cipher encryption
    key must be 16 or 24 bytes
    """
    category = "crypto (block)"
    name = "3des encrypt"
    icon = "wxART_CRYPT"

    def run(self, data:bytes, mode:("ecb","cbc","cfb")="ecb", key:bytes=b"0123456789abcdef", iv:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00", pad:bool=False):
        from Cryptodome.Cipher import DES3, DES
        from Cryptodome.Util import Padding
        if pad:
            data = Padding.pad(data, 8)
        if mode == "ecb":
            cryptor = DES3.new(key, DES3.MODE_ECB)
        elif mode == "cbc":
            cryptor = DES3.new(key, DES3.MODE_CBC, iv=iv)
        elif mode == "cfb":
            cryptor = DES3.new(key, DES3.MODE_CFB, iv=iv)
        return cryptor.encrypt(data)           



class ChaskeyLTSDecrypt(Transform):
    """
    Chaskey-LTS cipher decryption
    key must be 16 bytes

    Only supports CTR mode for now
    """
    category = "crypto (block)"
    name = "chaskey-lts decrypt"
    icon = "wxART_DECRYPT"

    def run(self, data:bytes, mode:("ctr",)="ctr", key:bytes=b"0123456789abcdef", counter:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"):
        from .libs.chaskey import Chaskey
        if len(counter) != 16:
            raise ValueError("Counter must be 16 bytes")
        elif len(key) != 16:
            raise ValueError("Key must be 16 bytes")
        return Chaskey(mode, key, counter).decrypt(data)


class ChaskeyLTSEncrypt(Transform):
    """
    Chaskey-LTS cipher encryption
    key must be 16 bytes

    Only supports CTR mode for now
    """
    category = "crypto (block)"
    name = "chaskey-lts encrypt"
    icon = "wxART_CRYPT"

    def run(self, data:bytes, mode:("ctr",)="ctr", key:bytes=b"0123456789abcdef", counter:bytes=b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"):
        from .libs.chaskey import Chaskey
        if len(counter) != 16:
            raise ValueError("Counter must be 16 bytes")
        elif len(key) != 16:
            raise ValueError("Key must be 16 bytes")
        return Chaskey(mode, key, counter).encrypt(data)    
