vkcoffee.py

from collections.abc import Generator
from typing import Optional, Tuple, Union
from base64 import encodebytes, decodebytes
from re import findall
from Crypto.Cipher import AES

class VKCoffeeCypher:
    DEFAULT_KEY = b"stupidUsersMustD"
    KEY_PADDING = b"mailRuMustDie"
    WRAPPERS = "(AP ID OG|PP|VK CO FF EE|VK C0 FF EE|II)"
    MSG_REGEX = WRAPPERS + r" ([A-F0-9\s]+) " + WRAPPERS

    def __init__(self, key: Optional[Union[int, bytes]] = DEFAULT_KEY):
        self.key = key

    @staticmethod
    def _pkcs7_pad(data: bytes) -> bytes:
        remainder = 16 - (len(data) % 16)
        if remainder == 16:
            return data
        return data + bytes([remainder] * remainder)

    @staticmethod
    def _pkcs7_unpad(data: bytes) -> bytes:
        if data[-1] > 16:
            return data
        return data[:-data[-1]]

    @classmethod
    def _encrypt_padded(cls, key: bytes, data: bytes) -> bytes:
        aes = AES.new(key, AES.MODE_ECB)
        return aes.encrypt(cls._pkcs7_pad(data))

    @classmethod
    def _decrypt_padded(cls, key: bytes, data: bytes) -> bytes:
        aes = AES.new(key, AES.MODE_ECB)
        return cls._pkcs7_unpad(aes.decrypt(data))

    @classmethod
    def extract_message(cls, text: str) -> Optional[str]:
        parts = findall(cls.MSG_REGEX, text)
        return parts[0][1] if parts else None

    @property
    def key(self):
        return self._key

    @key.setter
    def key(self, new_key: Optional[Union[int, str, bytes]] = None):
        if new_key is None:
            self._key = self.DEFAULT_KEY
        elif isinstance(new_key, bytes):
            self._key = self._pkcs7_pad(new_key)[:16]
        else:
            temp = str(new_key).encode('utf-8') + self.KEY_PADDING
            key = self._encrypt_padded(self.DEFAULT_KEY, temp)
            self._key = encodebytes(key).replace(b'\n', b'')[:16]

    def decrypt_raw(self, blob: str) -> bytes:
        return self._decrypt_padded(self._key, decodebytes(bytes.fromhex(blob)))

    def encrypt_raw(self, data: bytes) -> str:
        return str.join(" ", [
            "%02X" % c
            for c in encodebytes(self._encrypt_padded(self._key, data))
            if c != 0x0a
        ])

    def decrypt(self, message: str) -> Optional[str]:
        blob = self.extract_message(message)
        if not blob:
            return None
        return self.decrypt_raw(blob).decode('utf-8')

    def encrypt(self, message: str, wrapwith: str = "VK CO FF EE") -> str:
        return "{wrap} {data} {wrap}".format(
            wrap=wrapwith,
            data=self.encrypt_raw(message.encode("utf-8"))
        )

    @classmethod
    def bruteforce(cls, message: str) -> Generator[Tuple[int, str], None, None]:
        blob = cls.extract_message(message)
        if not blob:
            return
        for key in range(10000):
            try:
                res = cls(key).decrypt_raw(blob)
                if (res.strip()
                        and cls(key).encrypt_raw(res) == blob
                        and not set(res) & set(range(0x20))):
                    yield key, res.decode("utf-8")
            except Exception:
                pass

def test():
    from random import randint
    cof = VKCoffeeCypher(randint(0, 9999))
    print((data := cof.encrypt("Привет, мир!")))
    print(repr(cof.decrypt(data)))
    for k, text in VKCoffeeCypher.bruteforce(data):
        print('%4d' % k, repr(text))

if __name__ == "__main__":
    test()