import os
import socketserver
import textwrap

BLOCK_SIZE = 8
MOD_MASK = 0xFFFFFFFFFFFFFFFF

def generate_key() -> bytes:
    while True:
        key = os.urandom(16)
        a = int.from_bytes(key[:8], "little")
        if a & 1:
            return key

KEY = generate_key()
FLAG = os.getenv("FLAG", "vsosh{test_flag}")

def encrypt(plaintext: bytes) -> bytes:
    a = int.from_bytes(KEY[:8], "little")
    b = int.from_bytes(KEY[8:], "little")
    ciphertext = bytearray()
    for i in range(0, len(plaintext), BLOCK_SIZE):
        block = plaintext[i : i + BLOCK_SIZE].ljust(BLOCK_SIZE, b"\x00")
        p = int.from_bytes(block, "little")
        c = (a * p + b) & MOD_MASK
        ciphertext += c.to_bytes(BLOCK_SIZE, "little")
    return bytes(ciphertext)

BANNER = textwrap.dedent(
    """
    == Линейная ловушка ==
    Сервер шифрует флаг кастомным блочным шифром.
    Шифр для каждого блока: C = (A * P + B) mod 2^64
    A - нечётное 64-битное число, B - 64-битное число.
    Отправьте plaintext в hex (до 32 байт) или 'quit'.
    """
).strip()

PROMPT = "\n> "

class Handler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        conn = self.request
        conn.sendall((BANNER + "\n").encode())
        conn.sendall(("Encrypted flag (hex): " + encrypt(FLAG.encode()).hex() + "\n").encode())
        while True:
            conn.sendall(PROMPT.encode())
            line = self._readline(conn)
            if line is None:
                return
            line = line.strip()
            if not line:
                continue
            if line.lower() in {"quit", "exit"}:
                conn.sendall(b"Bye!\n")
                return
            if len(line) > 32 * 2:
                conn.sendall(b"Too long!\n")
                continue
            try:
                plaintext = bytes.fromhex(line)
            except ValueError:
                conn.sendall(b"Invalid hex!\n")
                continue
            ciphertext = encrypt(plaintext)
            conn.sendall(("Ciphertext (hex): " + ciphertext.hex() + "\n").encode())

    @staticmethod
    def _readline(conn, limit: int = 4096):
        data = bytearray()
        while len(data) < limit:
            chunk = conn.recv(1)
            if not chunk:
                return None
            if chunk == b"\n":
                break
            if chunk != b"\r":
                data += chunk
        return data.decode(errors="ignore")

class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True

if __name__ == "__main__":
    host = "0.0.0.0"
    port = 5000
    with ThreadedTCPServer((host, port), Handler) as server:
        server.serve_forever()
