from base64 import *
from Crypto.Random import get_random_bytes
from Crypto.Cipher import AES
from Crypto.Util.number import *
import os
import string
from hashlib import sha256
import socketserver
import signal
import random
# flag = os.getenv('FLAG').encode()

AES.block_size=16
iv = get_random_bytes(AES.block_size)
key = get_random_bytes(16)

class Task(socketserver.BaseRequestHandler):
    def proof_of_work(self):
        random.seed(os.urandom(8))
        proof = ''.join(
            [random.choice(string.ascii_letters+string.digits) for _ in range(20)])
        _hexdigest = sha256(proof.encode()).hexdigest()
        self.send(f"[+] sha256(XXXX+{proof[4:]}) == {_hexdigest}".encode())
        x = self.recv(prompt=b'[+] Plz tell me XXXX: ')
        if len(x) != 4 or sha256(x+proof[4:].encode()).hexdigest() != _hexdigest:
            return True
        return True

    def _recvall(self):
        BUFF_SIZE = 2048
        data = b''
        while True:
            part = self.request.recv(BUFF_SIZE)
            data += part
            if len(part) < BUFF_SIZE:
                break
        print("[+] "+ data.strip().decode())
        return data.strip()

    def send(self, msg, newline=True):
        try:
            if newline:
                msg += b'\n'
            self.request.sendall(msg)
            print('[-] '+msg.decode())
        except:
            pass

    def recv(self, prompt=b'> '):
        self.send(prompt, newline=False)
        return self._recvall()

    def timeout_handler(self, signum, frame):
        raise TimeoutError


class ThreadedServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    pass


def pad(text):
    pad_num = 16 - len(text) % 16
    return text + pad_num*bytes([pad_num])


def unpad(text):
    pad_num = int(text[-1])
    if pad_num == 0:
        return b'Man'
    for i in range(1, pad_num+1):
        if int(text[-i]) != pad_num:
            return b'Man'
    else:
        tmp = text[:-pad_num]
        return b'Mamba out!'


def encrypt(plain_text, key):
    cipher = AES.new(key, AES.MODE_CBC, iv)
    cipher_text = cipher.encrypt(pad(plain_text))
    return iv + cipher_text


def decrypt(cipher_text, key):
    iv = cipher_text[:AES.block_size]
    cipher = AES.new(key, AES.MODE_CBC, iv)
    tmp = cipher.decrypt(cipher_text[AES.block_size:])
    result = unpad(tmp)
    return result


class test(Task):
    def handle(self):
        if not self.proof_of_work():
            self.send(b'[!] Wrong!')
            return
        enc = encrypt(flag, key)
        self.send(b'secret:')
        self.send(b64encode(enc))
        self.send(b'What can I say?')
        while True:
            data = self.recv()
            try:
                self.send(decrypt(b64decode(data), key))
            except:
                self.send(b'Damn!')


if __name__ == "__main__":
    HOST, PORT = '0.0.0.0', 4444
    server = ThreadedServer((HOST, PORT), test)
    server.allow_reuse_address = True
    print(HOST, PORT)
    server.serve_forever()

from base64 import *
from pwn import *

# 工作量证明
def crack_proof_of_work(proof4:str,_hexdigest:str) -> bytes:
    from hashlib import sha256
    print("Task get: ",proof4,_hexdigest)
    import string
    alphabet = string.ascii_letters+string.digits
    for a in alphabet:
        for b in alphabet:
            for c in alphabet:
                for d in alphabet:
                    if sha256((a+b+c+d+proof4).encode()).hexdigest() == _hexdigest:
                        print(f"proof of work cracked: {a}{b}{c}{d}")
                        return (a+b+c+d).encode()

# 连接远程主机
s = remote("127.0.0.1", 4444)
proof_of_work_problem = s.recvline(keepends=False).decode()
print(s.recv()) # b'[+] Plz tell me XXXX: '
proof_of_work_answer = crack_proof_of_work(proof_of_work_problem[16:32],proof_of_work_problem[37:101])
s.sendline(proof_of_work_answer)
print(s.recvline(keepends=True)) # b'secret:\n'
encbase64bytes = s.recvline(keepends=False)
print(s.recv()) # b'What can I say?\n> '

flag = ''

for N in [0,16,32]: # 原题密文去掉IV向量就只有48字节分3段，按实际情况更改
    middle_list = [0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0]
    attack_list = [0,0,0,0, 0,0,0,0, 0,0,0,0, 0,0,0,0]
    attack_list.extend(list(b64decode(encbase64bytes))[16+N:32+N])
    s.send(b64encode(bytes(attack_list)))
    for index in range(15,-1,-1):
        for i in range(256):
            attack_list[index] = i
            s.send(b64encode(bytes(attack_list)))
            res = s.recvline(keepends=False)
            # print(res)
            _ = s.recv()
            if res != b'Man':
                break
        if index == 0:
            print(f"Got one bit! {index}",end=' \r')
            middle_list[index] = attack_list[index] ^ (16-index)
            break
        attack_list[index-1] ^= 0x0f # 反转比特 验证
        s.send(b64encode(bytes(attack_list)))
        res = s.recvline(keepends=False)
        # print(res)
        _ = s.recv()
        if res != b'Man':
            print(f"Got one bit! {index}",end='\r')
            middle_list[index] = attack_list[index] ^ (16-index)
            # 初始化下一padding
            for j in range(index,16):
                attack_list[j] = middle_list[j] ^ (16-index+1)
            index -= 1
        else:
            # 没想好怎么写这个分支，反正一般用不着吧（
            pass
    flag_ = bytes([middle_list[i] ^ list(b64decode(encbase64bytes))[0+N:16+N][i] for i in range(16)]).decode()
    flag += flag_
    print(f"\nflag got: {flag_} \n {flag}")