solve_home_roll_crypto.py

from sock import Sock
from binascii import hexlify, unhexlify
from itertools import product

class Cipher:
    BLOCK_SIZE = 16
    ROUNDS = 3
    def __init__(self, key):
        assert(len(key) == self.BLOCK_SIZE*self.ROUNDS)
        self.key = key

    def __block_encrypt(self, block):
        enc = int.from_bytes(block, "big")
        for i in range(self.ROUNDS):
            k = int.from_bytes(self.key[i*self.BLOCK_SIZE:(i+1)*self.BLOCK_SIZE], "big")
            enc &= k
            enc ^= k
        return hex(enc)[2:].rjust(self.BLOCK_SIZE*2, "0")


    def __pad(self, msg):
        if len(msg) % self.BLOCK_SIZE != 0:
            return msg + (bytes([0]) * (self.BLOCK_SIZE - (len(msg) % self.BLOCK_SIZE)))
        else:
            return msg
    
    def encrypt(self, msg):
        m = self.__pad(msg)
        e = ""
        for i in range(0, len(m), self.BLOCK_SIZE):
            e += self.__block_encrypt(m[i:i+self.BLOCK_SIZE])
        return e.encode()
        
def bruteforce(i: int):
    if i == 127:
        print(k1)
        return
    for m, n in product(range(2), range(2)):
        e1 = (m ^ 1) & n
        e2 = (int(k0[i]) | (m ^ 1)) & n
        if str(e1) == enc1[i] and str(e2) == enc2[i]:
            k1[i] = str(m)
            k2[i] = str(n)
            print(i, m, n)
            if i == 127:
                print(k1)
                print(k2)
                break
            else:
                bruteforce(i+1)
            k1[128-i-1] = 'x'
            k2[128-i-1] = 'x'
      

so = Sock('crypto.2021.chall.actf.co:21602')
k0 = ['x'] * 128
k1 = ['x'] * 128
k2 = ['x'] * 128
# Stage 1: send '111...1'
pt = 2**128-1
so.read_until(b"Would you like to encrypt [1], or try encrypting [2]? ")
so.send_line(b'1')
so.read_until(b"What would you like to encrypt: ")
so.send_line(hex(pt)[2:].zfill(16))
nk0_k2 = int(so.read_line().strip(), 16)

for i in range(128):
    so.read_until(b"Would you like to encrypt [1], or try encrypting [2]? ")
    so.send_line(b'1')
    so.read_until(b"What would you like to encrypt: ")
    so.send_line(hex(pt ^ (2**i))[2:].zfill(16))
    enc = int(so.read_line().strip(), 16)
    if enc == nk0_k2:
        k0[128-i-1] = '0'
    else:
        k0[128-i-1] = '1'
# Stage 2: send '000...000'

so.read_until(b"Would you like to encrypt [1], or try encrypting [2]? ")
so.send_line(b'1')
so.read_until(b"What would you like to encrypt: ")
so.send_line(b'0' * 32)
nk1_k2 = int(so.read_line().strip(), 16)
enc1 = bin(nk0_k2)[2:].zfill(16)
enc1=enc1.rjust(128,'0')
enc2 = bin(nk1_k2)[2:].zfill(16)
enc2=enc2.rjust(128,'0')
for i in range(128):
    for m, n in product(range(2), range(2)):
        e1 = (m ^ 1) & n
        e2 = (int(k0[i]) | (m ^ 1)) & n
        if str(e1) == enc1[i] and str(e2) == enc2[i]:
            k1[i] = str(m)
            k2[i] = str(n)
            break

#bruteforce(0)
key0 = int.to_bytes(int("".join(k0), 2), 16, 'big')
print(key0)
key1 = int.to_bytes(int("".join(k1), 2), 16, 'big')
print(key1)
key2 = int.to_bytes(int("".join(k2), 2), 16, 'big')
print(key2)
cipher = Cipher(key0 + key1 + key2)
so.read_until(b"Would you like to encrypt [1], or try encrypting [2]? ")
so.send_line(b'2')
for i in range(10):
	pt=so.read_line()
	pt=pt[14:-1]
	cipher = Cipher(key0 + key1 + key2)
	ciphertext = cipher.encrypt(unhexlify(pt)).decode()
	so.send_line(ciphertext)
	
tmp=so.read_line()
print(tmp)
flag=so.read_line()
print(flag)