#!/usr/bin/env python3
"""Local S7comm PLC simulator (server).

Stands in for a physical Siemens S7 PLC so that rtsp_multi_view can READ and
plc_sim can WRITE a small "process image" (byte array) over S7comm / ISO-on-TCP.

Handles: COTP Connect Request/Confirm, S7 Setup Communication, S7 Read Var
(function 0x04) and S7 Write Var (function 0x05). Multiple clients are served
on separate threads.

Usage:
    plc_server.py [-p PORT] [-l LEN] [--cycle] [-t CYCLE_MS]

    -p PORT    listen port (default 102; needs root on Linux)
    -l LEN     process image size in bytes (default 16)
    --cycle    also rotate the 6 layout-switch states every cycle
    -t MS      state cycle interval in ms (default 10000)

With --cycle the server drives byte0 through 0x00,0x03,0x0C,0x30,0x0F,0x3F
all by itself, so plc_sim is not required.
"""

import argparse
import socket
import sys
import threading
import time

TS_BYTE = 0x02
CYCLE_STATES = [0x00, 0x03, 0x0C, 0x30, 0x0F, 0x3F]
STATE_NAMES = [
    "state 0: all bits 0 (default layout)",
    "state 1: M0.0|M0.1 (layout 0)",
    "state 2: M0.2|M0.3 (layout 1)",
    "state 3: M0.4|M0.5 (layout 2)",
    "state 4: M0.0,0.1;M0.2,0.3 (layout 3)",
    "state 5: M0.0,0.1;M0.2,0.3;M0.4,0.5 (layout 4)",
]

CC = bytes([0x03, 0x00, 0x00, 0x15]) + bytes([
    0x0B, 0xD0, 0x00, 0x00, 0x01, 0x00,
    0xC0, 0x01, 0x0A,
    0xC1, 0x02, 0x01, 0x00,
    0xC2, 0x02, 0x01, 0x02,
])


def put16(v):
    return bytes([(v >> 8) & 0xFF, v & 0xFF])


def tpkt(s7):
    total = 4 + 3 + len(s7)
    return bytes([0x03, 0x00, (total >> 8) & 0xFF, total & 0xFF,
                  0x02, 0xF0, 0x80]) + s7


def recv_exact(c, n):
    buf = b''
    while len(buf) < n:
        r = c.recv(n - len(buf))
        if not r:
            return None
        buf += r
    return buf


def read_tpkt(c):
    h = recv_exact(c, 4)
    if not h:
        return None
    total = (h[2] << 8) | h[3]
    if total < 7 or total > 4096:
        return None
    body = recv_exact(c, total - 4)
    if body is None:
        return None
    return h + body


class PlcServer:
    def __init__(self, area_len, cycle, cycle_ms):
        self.lock = threading.Lock()
        self.area = bytearray(area_len)
        self.cycle = cycle
        self.cycle_ms = cycle_ms
        self.state_idx = 0

    def setup_resp(self, seq):
        return bytes([0x32, 0x03, 0x00, 0x00]) + seq + \
            bytes([0x00, 0x08, 0x00, 0x00, 0x00, 0x00,
                   0xF0, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0xE0])

    def read_var(self, seq, s7):
        ts = s7[15]
        n = (s7[16] << 8) | s7[17]
        if n > len(self.area):
            n = len(self.area)
        with self.lock:
            data = bytes(self.area[:n])
        blen = n if ts == TS_BYTE else n * 8
        resp = bytes([0x32, 0x03, 0x00, 0x00]) + seq + \
            put16(2) + put16(4 + n) + bytes([0x00, 0x00])
        resp += bytes([0x04, 0x01])
        resp += bytes([0xFF, ts]) + put16(blen) + data
        return resp

    def write_var(self, seq, s7):
        ts = s7[15]
        n = (s7[16] << 8) | s7[17]
        item = s7[24:]
        blen = (item[2] << 8) | item[3]
        nbytes = blen if ts == TS_BYTE else blen // 8
        if nbytes > len(self.area):
            nbytes = len(self.area)
        data = item[4:4 + nbytes]
        with self.lock:
            self.area[:nbytes] = data
        print("[plc] write %d bytes @%d -> byte0=0x%02X" %
              (nbytes, 0, data[0] if data else 0), flush=True)
        return bytes([0x32, 0x03, 0x00, 0x00]) + seq + \
            bytes([0x00, 0x0E, 0x00, 0x00, 0x00, 0x00, 0x05, 0x01, 0xFF])

    def handle(self, conn, addr):
        try:
            conn.settimeout(120)
            while True:
                frame = read_tpkt(conn)
                if frame is None:
                    break
                payload = frame[4:]
                if payload[1] == 0xE0:            # COTP Connect Request
                    conn.sendall(CC)
                elif payload[1] == 0xF0:          # COTP Data
                    s7 = payload[3:]
                    if len(s7) < 12:
                        continue
                    func = s7[10]
                    seq = s7[4:6]
                    if func == 0xF0:              # S7 Setup Communication
                        conn.sendall(tpkt(self.setup_resp(seq)))
                    elif func == 0x04:            # S7 Read Var
                        conn.sendall(tpkt(self.read_var(seq, s7)))
                    elif func == 0x05:            # S7 Write Var
                        conn.sendall(tpkt(self.write_var(seq, s7)))
                    # else: ignore unknown requests
        except Exception:
            pass
        finally:
            conn.close()
        print("[plc] client %s:%d disconnected" % addr, flush=True)

    def cycle_loop(self):
        while True:
            time.sleep(self.cycle_ms / 1000.0)
            v = CYCLE_STATES[self.state_idx]
            with self.lock:
                self.area[0] = v
            print("[plc] %s (byte0=0x%02X)" %
                  (STATE_NAMES[self.state_idx], v), flush=True)
            self.state_idx = (self.state_idx + 1) % len(CYCLE_STATES)


def main():
    ap = argparse.ArgumentParser(description="Local S7comm PLC simulator")
    ap.add_argument("-p", "--port", type=int, default=102)
    ap.add_argument("-l", "--len", type=int, default=16)
    ap.add_argument("--cycle", action="store_true",
                    help="self-rotate the 6 layout states every cycle")
    ap.add_argument("-t", "--cycle-ms", type=int, default=10000)
    args = ap.parse_args()

    server = PlcServer(args.len, args.cycle, args.cycle_ms)

    if args.cycle:
        threading.Thread(target=server.cycle_loop, daemon=True).start()

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("0.0.0.0", args.port))
    srv.listen(8)
    print("[plc] local PLC listening on 0.0.0.0:%d area_len=%d%s" %
          (args.port, args.len, " (cycling)" if args.cycle else ""), flush=True)

    while True:
        conn, addr = srv.accept()
        print("[plc] client %s:%d connected" % addr, flush=True)
        threading.Thread(target=server.handle, args=(conn, addr),
                         daemon=True).start()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)