from cryptography.hazmat.primitives.ciphers.algorithms import AES, SM4
from cryptography.hazmat.primitives.ciphers import Cipher, modes
import os


class Paranoia:
    def __init__(self, keys):
        self.keys = keys

    def __pad(self, data: bytes, bs: int) -> bytes:
        return data + (chr(bs - len(data) % bs) * (bs - len(data) % bs)).encode()

    def __encrypt(self, algorithm, data: bytes, key: bytes):
        cipher = Cipher(algorithm(key), modes.ECB())
        encryptor = cipher.encryptor()
        return encryptor.update(data) + encryptor.finalize()

    def encrypt(self, data: bytes):
        """
        🇨🇳 encryption to protect against the 🇺🇸 backdoor and
        🇺🇸 encryption to protect against the 🇨🇳 backdoor

        I'm a genius !
        """

        data = self.__pad(data, 16)
        data = self.__encrypt(AES, data, self.keys[0])
        data = self.__encrypt(SM4, data, self.keys[1])
        return data


with open("flag.txt", "rb") as f:
    flag = f.read()

keys = [os.urandom(16) for _ in range(2)]
paranoia = Paranoia(keys)

banner = b"I don't trust governments, thankfully I've found smart a way to keep my data secure."

print("pt_banner =", banner)
print("ct_banner =", paranoia.encrypt(banner))
print("enc_flag  =", paranoia.encrypt(flag))

# To comply with cryptography export regulations,
# 6 bytes = 2**48 bits, should be bruteforce-proof anyway
for n, k in enumerate(keys):
    print(f"k{n} = {k[3:]}")

pt_banner = b"I don't trust governments, thankfully I've found smart a way to keep my data secure."
ct_banner = b"\xd5\xae\x14\x9de\x86\x15\x88\xe0\xdc\xc7\x88{\xcfy\x81\x91\xbaH\xb6\x06\x02\xbey_0\xa5\x8a\xf6\x8b?\x9c\xc9\x92\xac\xdeb=@\x9bI\xeeY\xa0\x8d/o\xfa%)\xfb\xa2j\xd9N\xf7\xfd\xf6\xc2\x0b\xc3\xd2\xfc\te\x99\x9aIG\x01_\xb3\xf4\x0fG\xfb\x9f\xab\\\xe0\xcc\x92\xf5\xaf\xa2\xe6\xb0h\x7f}\x92O\xa6\x04\x92\x88"
enc_flag = b"\xaf\xe0\xb8h=_\xb0\xfbJ0\xe6l\x8c\xf2\xad\x14\xee\xccw\xe9\xff\xaa\xb2\xe9c\xa4\xa0\x95\x81\xb8\x03\x93\x7fg\x00v\xde\xba\xfe\xb92\x04\xed\xc4\xc7\x08\x8c\x96C\x97\x07\x1b\xe8~':\x91\x08\xcf\x9e\x81\x0b\x9b\x15"
k0 = b'C\xb0\xc0f\xf3\xa8\n\xff\x8e\x96g\x03"'
k1 = b"Q\x95\x8b@\xfbf\xba_\x9e\x84\xba\x1a7"

from cryptography.hazmat.primitives.ciphers.algorithms import AES, SM4
from cryptography.hazmat.primitives.ciphers import Cipher, modes
import os

pt_banner = b"I don't trust governments, thankfully I've found smart a way to keep my data secure."
ct_banner = b"\xd5\xae\x14\x9de\x86\x15\x88\xe0\xdc\xc7\x88{\xcfy\x81\x91\xbaH\xb6\x06\x02\xbey_0\xa5\x8a\xf6\x8b?\x9c\xc9\x92\xac\xdeb=@\x9bI\xeeY\xa0\x8d/o\xfa%)\xfb\xa2j\xd9N\xf7\xfd\xf6\xc2\x0b\xc3\xd2\xfc\te\x99\x9aIG\x01_\xb3\xf4\x0fG\xfb\x9f\xab\\\xe0\xcc\x92\xf5\xaf\xa2\xe6\xb0h\x7f}\x92O\xa6\x04\x92\x88"
enc_flag = b"\xaf\xe0\xb8h=_\xb0\xfbJ0\xe6l\x8c\xf2\xad\x14\xee\xccw\xe9\xff\xaa\xb2\xe9c\xa4\xa0\x95\x81\xb8\x03\x93\x7fg\x00v\xde\xba\xfe\xb92\x04\xed\xc4\xc7\x08\x8c\x96C\x97\x07\x1b\xe8~':\x91\x08\xcf\x9e\x81\x0b\x9b\x15"
k0 = b'C\xb0\xc0f\xf3\xa8\n\xff\x8e\x96g\x03"'
k1 = b"Q\x95\x8b@\xfbf\xba_\x9e\x84\xba\x1a7"

def pad(data: bytes, bs: int) -> bytes:
    return data + (chr(bs - len(data) % bs) * (bs - len(data) % bs)).encode()

def aes_encrypt(key, data):
    cipher = Cipher(AES(key), modes.ECB())
    encryptor = cipher.encryptor()
    return encryptor.update(data) + encryptor.finalize()

def sm4_decrypt(key, data):
    cipher = Cipher(SM4(key), modes.ECB())
    decryptor = cipher.decryptor()
    return decryptor.update(data) + decryptor.finalize()

# Take only the first block of the banner to speed-up
banner = pad(b"I don't trust go", 16)
# Take only the first block of the encrypted banner to speed-up
encrypted_banner = ct_banner[:16]

partial_encrypted_banner = {}
partial_decrypted_banner = {}

for i in range(2**24):

    # Get the key for SM4 decryption and AES encryption
    key = hex(i)[2:]
    key = bytes.fromhex("0"*(6-len(key)) + key)
    test_key0 = key + k0
    test_key1 = key + k1

    # Compute the AES encryption and the SM4 decryption using the guessed keys
    aes_e = aes_encrypt(test_key0, banner)[:16].hex()
    sm4_d = sm4_decrypt(test_key1, encrypted_banner)[:16].hex()
    # Store them
    partial_encrypted_banner[aes_e] = test_key0
    partial_decrypted_banner[sm4_d] = test_key1

# Compute the intersection of the sets
print(partial_encrypted_banner.keys() & partial_encrypted_banner.keys())

# We got the keys from the intersection
aes_key = b'If-C\xb0\xc0f\xf3\xa8\n\xff\x8e\x96g\x03"'
sm4_key = b'\x94\xcb\x92Q\x95\x8b@\xfbf\xba_\x9e\x84\xba\x1a7'
decryptor_aes = Cipher(AES(aes_key), modes.ECB()).decryptor()
decryptor_sm4 = Cipher(SM4(sm4_key), modes.ECB()).decryptor()

# Flag decryption
print(
    decryptor_aes.update(decryptor_sm4.update(enc_flag) + decryptor_sm4.finalize()) + decryptor_aes.finalize()
)