from random import choice as c
from random import randint, shuffle
import socketserver
import signal
import string
import random
import os

class Task(socketserver.BaseRequestHandler):
    def _recvall(self):
        BUFF_SIZE = 2048
        data = b''
        while True:
            part = self.request.recv(BUFF_SIZE)
            data += part
            if len(part) < BUFF_SIZE:
                break
        return data.strip()

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

    def recv(self, prompt=b'[-] '):
        self.send(prompt, newline=False)
        return self._recvall()
    
    # 定义一个函数 calc，用于计算问题的结果
    def calc(self, ans, chests, expr):
        try:
            C0, C1, C2, C3, C4, C5, C6 = chests
            r = eval(expr)
        except Exception as e:
            self.send(b"Shannon fails to understand your words.\n")
            self.send(e)
            exit(0)
        return ans(r)
    
    # 定义一个游戏回合函数 do_round
    def do_round(self):
        # 定义两个函数 truth 和 lie，用于判断 Shannon 的回答是否为真或假
        truth = lambda r: not not r
        lie = lambda r: not r
        chests = []
        for i in range(7):
            # 随机生成 7 个 True 或 False，表示每个箱子是宝藏还是怪物
            chests.append(c((True, False)))
        self.send(b"Seven chests lie here, with mimics or treasure hidden inside.\nBut don't worry. Trusty Shannon knows what to do.")
        # 随机选择 Shannon 的回答策略，包括多少次 lie
        lie_count = c((1, 2))
        Shannon = [truth] * (15 - lie_count) + [lie] * lie_count
        # 猜猜第几次是在说谎 :-)
        shuffle(Shannon)
        for i in range(15):
            self.send(b"Ask Shannon:")
            question = self.recv().decode().strip()  # 输入问题
            for word in question.split(" "):
                if word not in white_list:  # 检查问题中是否包含白名单以外的关键词或符号
                    self.send("({}) No treasure for dirty hacker!".format(word).encode())
                    exit(0)
            res = str(self.calc(Shannon[i], chests, question))  # 计算问题的结果
            self.send('Shannon answers: {}!\n'.format(res).encode())
        self.send(b"Now open the chests:")
        # 输入0或1表示每个宝箱的真假，用空格隔开
        chests_string = self.recv().decode()
        return chests == list(map(int, chests_string.strip().split(" ")))

    def handle(self):
        signal.alarm(60)

        # 游戏开始
        self.send(b"The Unbreakable Shannon has returned, with some suspicious chests and a far more complicated strategy -- he MAY LIE ONCE OR TWICE! Can you still get all the treasure without losing your head?")
        for i in range(1):
            if not self.do_round():  # 执行游戏回合
                self.send(b"A chest suddenly comes alive and BITE YOUR HEAD OFF.\n")
                return
        self.send("You've found all the treasure! {}\n".format(flag).encode())  # 赢得游戏，获得flag


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


class ForkedServer(socketserver.ForkingMixIn, socketserver.TCPServer):
    pass


if __name__ == "__main__":
    flag = bytes(os.getenv("FLAG"),"utf-8") # type: ignore
    # 白名单列表，包含允许在问题中使用的关键词和符号
    white_list = ['==', '(', ')', 'C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', '0', '1', 'and', 'or', 'not']
    HOST, PORT = '0.0.0.0', 10011
    server = ForkedServer((HOST, PORT), Task)
    server.allow_reuse_address = True
    print(HOST, PORT)
    server.serve_forever()