#!/usr/bin/env python3
"""Мок-модуль кондиционера (сторона устройства) для интеграционных тестов
сессии. stdlib-only: AES-256 реализован на чистом python (объёмы крошечные).

Сценарные флаги:
  --503               всегда отвечать 503 на local_reg (нет слотов)
  --no-poll           key exchange без опроса commands.json (зависание)
  --rekey-every N     ре-кей на каждый N-й local_reg (N=1 — каждый)
  --stale-gap S       ре-кей, если зазор между local_reg >= S секунд
                      (эмуляция «вернувшегося» приложения; по умолчанию 44)
  --garbage-pushes N  первые N push с отрезанным блоком шифротекста
                      (входящая цепочка расходится на 1 сообщение)
  --break-outbound N  N раз «не заметить» ответ commands.json (исходящий
                      десинк): затем подпись наших команд не сойдётся —
                      мок, как реальный модуль, ре-кает на следующем local_reg
  --fail-pushes N     первые N push'ей с испорченной подписью
  --push-every S      спонтанный push свойства tick каждые S секунд
  --fail-first-ke     первый key_exchange с ver=2 (ожидаем 426)
Вывод (stdout, строки):
  REG <first|put> notify=0|1
  KE <random1>
  CMD <method> <resource|name> [value]
  PUSH <name> <value>
  DELETE
"""
import argparse
import base64
import hashlib
import hmac
import http.server
import json
import random
import socket
import string
import sys
import threading
import time

# ---------------------------------------------------------------------------
# Чистый python AES-256 (encrypt/decrypt block), CBC поверх.
# ---------------------------------------------------------------------------
_SBOX = [
    0x63,0x7c,0x77,0x7b,0xf2,0x6b,0x6f,0xc5,0x30,0x01,0x67,0x2b,0xfe,0xd7,0xab,0x76,
    0xca,0x82,0xc9,0x7d,0xfa,0x59,0x47,0xf0,0xad,0xd4,0xa2,0xaf,0x9c,0xa4,0x72,0xc0,
    0xb7,0xfd,0x93,0x26,0x36,0x3f,0xf7,0xcc,0x34,0xa5,0xe5,0xf1,0x71,0xd8,0x31,0x15,
    0x04,0xc7,0x23,0xc3,0x18,0x96,0x05,0x9a,0x07,0x12,0x80,0xe2,0xeb,0x27,0xb2,0x75,
    0x09,0x83,0x2c,0x1a,0x1b,0x6e,0x5a,0xa0,0x52,0x3b,0xd6,0xb3,0x29,0xe3,0x2f,0x84,
    0x53,0xd1,0x00,0xed,0x20,0xfc,0xb1,0x5b,0x6a,0xcb,0xbe,0x39,0x4a,0x4c,0x58,0xcf,
    0xd0,0xef,0xaa,0xfb,0x43,0x4d,0x33,0x85,0x45,0xf9,0x02,0x7f,0x50,0x3c,0x9f,0xa8,
    0x51,0xa3,0x40,0x8f,0x92,0x9d,0x38,0xf5,0xbc,0xb6,0xda,0x21,0x10,0xff,0xf3,0xd2,
    0xcd,0x0c,0x13,0xec,0x5f,0x97,0x44,0x17,0xc4,0xa7,0x7e,0x3d,0x64,0x5d,0x19,0x73,
    0x60,0x81,0x4f,0xdc,0x22,0x2a,0x90,0x88,0x46,0xee,0xb8,0x14,0xde,0x5e,0x0b,0xdb,
    0xe0,0x32,0x3a,0x0a,0x49,0x06,0x24,0x5c,0xc2,0xd3,0xac,0x62,0x91,0x95,0xe4,0x79,
    0xe7,0xc8,0x37,0x6d,0x8d,0xd5,0x4e,0xa9,0x6c,0x56,0xf4,0xea,0x65,0x7a,0xae,0x08,
    0xba,0x78,0x25,0x2e,0x1c,0xa6,0xb4,0xc6,0xe8,0xdd,0x74,0x1f,0x4b,0xbd,0x8b,0x8a,
    0x70,0x3e,0xb5,0x66,0x48,0x03,0xf6,0x0e,0x61,0x35,0x57,0xb9,0x86,0xc1,0x1d,0x9e,
    0xe1,0xf8,0x98,0x11,0x69,0xd9,0x8e,0x94,0x9b,0x1e,0x87,0xe9,0xce,0x55,0x28,0xdf,
    0x8c,0xa1,0x89,0x0d,0xbf,0xe6,0x42,0x68,0x41,0x99,0x2d,0x0f,0xb0,0x54,0xbb,0x16]
_RCON = [0x01,0x02,0x04,0x08,0x10,0x20,0x40,0x80,0x1b,0x36,0x6c,0xd8,0xab,0x4d]
_INV_SBOX = [0]*256
for _i, _b in enumerate(_SBOX):
    _INV_SBOX[_b] = _i

def _xtime(a):
    a <<= 1
    if a & 0x100:
        a = (a ^ 0x1b) & 0xff
    return a

def _expand_key(key):  # 32 байта -> 60 слов по 4 байта (flat список)
    w = list(key)
    for i in range(32, 240, 4):
        t = w[i-4:i]
        if i % 32 == 0:
            t = t[1:] + t[:1]
            t = [_SBOX[b] for b in t]
            t[0] ^= _RCON[i//32 - 1]
        elif i % 32 == 16:
            t = [_SBOX[b] for b in t]
        w += [w[i-32+j] ^ t[j] for j in range(4)]
    return w

def _aes_encrypt_block(w, block):
    s = list(block)
    def add_round_key(r):
        for i in range(16):
            s[i] ^= w[r*16 + i]
    def sub_shift():
        # SubBytes + ShiftRows (строка r — байты r, r+4, r+8, r+12 — влево на r)
        t = [_SBOX[b] for b in s]
        out = [0]*16
        for r in range(4):
            for c in range(4):
                out[r + 4*c] = t[r + 4*((c + r) % 4)]
        for i in range(16):
            s[i] = out[i]
    def mix():
        t = [0]*16
        for c in range(4):
            col = s[c*4:c*4+4]
            t[c*4+0] = _xtime(col[0]) ^ _xtime(col[1]) ^ col[1] ^ col[2] ^ col[3]
            t[c*4+1] = col[0] ^ _xtime(col[1]) ^ _xtime(col[2]) ^ col[2] ^ col[3]
            t[c*4+2] = col[0] ^ col[1] ^ _xtime(col[2]) ^ _xtime(col[3]) ^ col[3]
            t[c*4+3] = _xtime(col[0]) ^ col[0] ^ col[1] ^ col[2] ^ _xtime(col[3])
        for i in range(16):
            s[i] = t[i]
    add_round_key(0)
    for rnd in range(1, 14):
        sub_shift(); mix(); add_round_key(rnd)
    sub_shift(); add_round_key(14)
    return bytes(s)

def _aes_decrypt_block(w, block):
    inv_sbox = _INV_SBOX
    def inv_sub_shift(s):
        t = [inv_sbox[b] for b in s]
        out = [0]*16
        for r in range(4):
            for c in range(4):
                out[r + 4*c] = t[r + 4*((c - r) % 4)]  # инверсия сдвига влево на r
        return out
    def inv_mix(s):
        def mul(a, b):
            p = 0
            for _ in range(8):
                if b & 1:
                    p ^= a
                hi = a & 0x80
                a = (a << 1) & 0xff
                if hi:
                    a ^= 0x1b
                b >>= 1
            return p
        t = [0]*16
        for c in range(4):
            col = s[c*4:c*4+4]
            t[c*4+0] = mul(col[0],14) ^ mul(col[1],11) ^ mul(col[2],13) ^ mul(col[3],9)
            t[c*4+1] = mul(col[0],9) ^ mul(col[1],14) ^ mul(col[2],11) ^ mul(col[3],13)
            t[c*4+2] = mul(col[0],13) ^ mul(col[1],9) ^ mul(col[2],14) ^ mul(col[3],11)
            t[c*4+3] = mul(col[0],11) ^ mul(col[1],13) ^ mul(col[2],9) ^ mul(col[3],14)
        return t
    s = list(block)
    def add_round_key(r):
        for i in range(16):
            s[i] ^= w[r*16 + i]
    add_round_key(14)
    for rnd in range(13, 0, -1):
        s = inv_sub_shift(s)   # InvShiftRows + InvSubBytes (коммутируют)
        add_round_key(rnd)
        s = inv_mix(s)
    s = inv_sub_shift(s)
    add_round_key(0)
    return bytes(s)

class PyAes:
    def __init__(self, key):
        self.w = _expand_key(key)
    def cbc_encrypt(self, iv, data):
        out = b""
        prev = iv
        for i in range(0, len(data), 16):
            blk = data[i:i+16]
            blk = bytes(a ^ b for a, b in zip(blk, prev))
            prev = _aes_encrypt_block(self.w, blk)
            out += prev
        return out, prev
    def cbc_decrypt(self, iv, data):
        out = b""
        prev = iv
        for i in range(0, len(data), 16):
            blk = data[i:i+16]
            dec = _aes_decrypt_block(self.w, blk)
            out += bytes(a ^ b for a, b in zip(dec, prev))
            prev = blk
        return out, prev

# ---------------------------------------------------------------------------
class MockCrypto:
    """Ключи одной стороны мока: dev (исходящие push) и app (входящие команды)."""
    def __init__(self, lanip_key, rnd1, rnd2, t1, t2):
        k = lanip_key.encode()
        b1, b2 = rnd1.encode(), rnd2.encode()
        s1, s2 = str(t1).encode(), str(t2).encode()
        def m(msg, suf):
            msg = msg + bytes([suf])
            return hmac.digest(k, hmac.digest(k, msg, "sha256") + msg, "sha256")
        A, D = b1 + b2 + s1 + s2, b2 + b1 + s2 + s1
        self.dev_sign, self.dev_aes = m(D, 0x30), PyAes(m(D, 0x31))
        self.app_sign, self.app_aes = m(A, 0x30), PyAes(m(A, 0x31))
        self.dev_iv, self.app_iv = m(D, 0x32)[:16], m(A, 0x32)[:16]

    def pack_push(self, seq, data_json):
        plain = json.dumps({"seq_no": seq, "data": data_json},
                           separators=(",", ":")).encode()
        sign = base64.b64encode(hmac.digest(self.dev_sign, plain, "sha256")).decode()
        n = ((len(plain) + 1 + 15) // 16) * 16
        ct, self.dev_iv = self.dev_aes.cbc_encrypt(self.dev_iv, plain.ljust(n, b"\x00"))
        enc = base64.b64encode(ct).decode()
        return json.dumps({"enc": enc, "sign": sign}, separators=(",", ":"))

    def unpack_command(self, body):
        d = json.loads(body)
        pt, self.app_iv = self.app_aes.cbc_decrypt(
            self.app_iv, base64.b64decode(d["enc"]))
        pt = pt.rstrip(b"\x00")
        # Подпись проверяется ДО разбора JSON (мусор не парсим).
        ok = base64.b64encode(hmac.digest(self.app_sign, pt, "sha256")).decode() == d["sign"]
        if not ok:
            return False, None
        return True, json.loads(pt.decode())

# ---------------------------------------------------------------------------
class Mock:
    def __init__(self, args):
        self.args = args
        self.lock = threading.Lock()
        self.props = {"operation_mode": 6, "fan_speed": 4, "tick": 0}
        self.crypto = None
        self.app_addr = None      # (ip, port) приложения
        self.push_seq = 0
        self.reg_count = 0
        self.last_reg_time = None
        self.fail_pushes = args.fail_pushes
        self.garbage_pushes = args.garbage_pushes
        self.miss_response = args.break_outbound
        self.stop = threading.Event()

    # ---------- исходящие к приложению ----------
    def http_call(self, method, path, body=b"", timeout=5):
        ip, port = self.app_addr
        c = socket.create_connection((ip, port), timeout=timeout)
        req = (f"{method} {path} HTTP/1.1\r\nHost: {ip}\r\n"
               f"Content-Type: application/json\r\n"
               f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n").encode() + body
        c.sendall(req)
        raw = b""
        while True:
            chunk = c.recv(4096)
            if not chunk:
                break
            raw += chunk
        c.close()
        head, _, resp_body = raw.partition(b"\r\n\r\n")
        status = int(head.split(b" ")[1])
        return status, resp_body

    def do_key_exchange(self):
        rnd1 = "".join(random.choice(string.ascii_letters + string.digits + "+/")
                       for _ in range(16))
        t1 = int(time.monotonic_ns() // 1000)
        ver = 2 if (self.args.fail_first_ke and self.reg_count == 1) else 1
        body = json.dumps({"key_exchange": {
            "ver": ver, "proto": 1, "key_id": self.args.key_id,
            "random_1": rnd1, "time_1": t1, "sec": ""}},
            separators=(",", ":")).encode()
        status, resp = self.http_call("POST", "/local_lan/key_exchange.json", body)
        print(f"KE {rnd1} -> {status}", flush=True)
        if status != 200:
            return False
        d = json.loads(resp)
        self.crypto = MockCrypto(self.args.lanip_key, rnd1, d["random_2"],
                                 t1, d["time_2"])
        return True

    def poll_commands(self):
        """Опрашивает commands.json, пока 206; исполняет команды."""
        while True:
            status, body = self.http_call("GET", "/local_lan/commands.json")
            if status != 200 and status != 206:
                print(f"POLL -> {status}", flush=True)
                return
            if not self.crypto:
                return
            if self.miss_response > 0:
                # «Модуль не получил/не расшифровал ответ»: цепочка приложения
                # ушла, у мока нет — исходящий десинк.
                self.miss_response -= 1
                print("CMD skipped (outbound desync)", flush=True)
                return
            ok, payload = self.crypto.unpack_command(body)
            if not ok:
                print("CMD bad-sign", flush=True)
                # Реальный модуль: на следующем local_reg — key exchange.
                self.crypto = None
                return
            data = payload.get("data", {})
            cmds = data.get("cmds", [])
            props = data.get("properties", [])
            if cmds:
                for c in cmds:
                    cmd = c.get("cmd", {})
                    if cmd.get("method") == "DELETE":
                        print("DELETE", flush=True)
                        return
                    res = cmd.get("resource", "")
                    name = res.split("name=")[-1]
                    cid = cmd.get("cmd_id", -1)
                    print(f"CMD GET {name} cid={cid}", flush=True)
                    self.push_datapoint(name, cid=cid)
            elif props:
                for p in props:
                    pr = p.get("property", {})
                    self.props[pr.get("name", "?")] = pr.get("value")
                    print(f"CMD SET {pr.get('name')}={pr.get('value')}", flush=True)
            elif not data:
                return
            if status == 200:
                return

    def push_datapoint(self, name, cid=-1, corrupt=False, garbage=False):
        self.push_seq += 1
        body = self.crypto.pack_push(self.push_seq - 1,
                                     {"name": name, "value": self.props.get(name, 0)})
        if garbage:
            # РЕАЛЬНЫЙ десинк CBC: отрезать последний блок шифротекста —
            # цепочка мока ушла на блок дальше, приложение отстанет.
            d = json.loads(body)
            ct = base64.b64decode(d["enc"])
            d["enc"] = base64.b64encode(ct[:-16]).decode()
            body = json.dumps(d, separators=(",", ":"))
        if corrupt:
            d = json.loads(body)
            s = d["sign"]
            d["sign"] = ("A" if s[-2] != "A" else "B") + s[1:]
            body = json.dumps(d, separators=(",", ":"))
        path = "/local_lan/property/datapoint.json"
        if cid >= 0:
            path += f"?cmd_id={cid}&status=200"
        status, _ = self.http_call("POST", path, body.encode())
        print(f"PUSH {name} -> {status}{' CORRUPT' if corrupt else ''}", flush=True)

    # ---------- входящие local_reg ----------
    def handle_local_reg(self, body):
        self.reg_count += 1
        d = json.loads(body)["local_reg"]
        notify = d.get("notify", 0)
        self.app_addr_json = d
        self.app_addr = (d["ip"], d["port"])
        print(f"REG {'first' if self.reg_count == 1 else 'put'} notify={notify}",
              flush=True)
        now = time.monotonic()
        stale = (self.last_reg_time is not None and
                 now - self.last_reg_time >= self.args.stale_gap)
        self.last_reg_time = now
        if self.args.http503:
            return 503
        need_ke = self.crypto is None or stale or (
            self.args.rekey_every and self.reg_count % self.args.rekey_every == 0)
        if need_ke and not self.do_key_exchange():
            return 202
        if not self.args.no_poll:
            self.poll_commands()
        return 202

    def spontaneous_loop(self):
        while not self.stop.wait(self.args.push_every or 10):
            if self.crypto and not self.args.no_poll:
                with self.lock:
                    self.props["tick"] += 1
                    corrupt = self.fail_pushes > 0
                    if corrupt:
                        self.fail_pushes -= 1
                    garbage = self.garbage_pushes > 0
                    if garbage:
                        self.garbage_pushes -= 1
                    self.push_datapoint("tick", corrupt=corrupt, garbage=garbage)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, required=True)
    ap.add_argument("--lanip-key", required=True)
    ap.add_argument("--key-id", type=int, required=True)
    ap.add_argument("--503", dest="http503", action="store_true")
    ap.add_argument("--no-poll", action="store_true")
    ap.add_argument("--rekey-every", type=int, default=0)
    ap.add_argument("--fail-pushes", type=int, default=0)
    ap.add_argument("--garbage-pushes", type=int, default=0)
    ap.add_argument("--break-outbound", type=int, default=0)
    ap.add_argument("--stale-gap", type=float, default=44.0)
    ap.add_argument("--push-every", type=float, default=0)
    ap.add_argument("--fail-first-ke", action="store_true")
    args = ap.parse_args()

    mock = Mock(args)
    lock = mock.lock

    class H(http.server.BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"
        def log_message(self, *a):
            pass
        def do_POST(self):
            n = int(self.headers.get("Content-Length") or 0)
            body = self.rfile.read(n)
            if self.path.startswith("/local_reg.json"):
                with lock:
                    code = mock.handle_local_reg(body)
                self.send_response(code)
                self.send_header("Content-Length", "0")
                self.end_headers()
                return
            self.send_response(404)
            self.send_header("Content-Length", "0")
            self.end_headers()
        def do_PUT(self):
            self.do_POST()

    srv = http.server.ThreadingHTTPServer(("127.0.0.1", args.port), H)
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    if args.push_every:
        threading.Thread(target=mock.spontaneous_loop, daemon=True).start()
    print("READY", flush=True)
    try:
        while True:
            time.sleep(0.5)
    except KeyboardInterrupt:
        pass

if __name__ == "__main__":
    main()