#!/usr/bin/env python3
"""
inject_deformed_packet.py — Modbus 畸形包注入测试工具
=====================================================
用途:     通过 485 穿透链路，向光伏跟踪器主控板发送压力测试报文
目标:     验证 STM32 从机在恶劣电磁环境下的防死机/防崩溃能力

三种攻击模式:
  1. CRC 错误包     → 检验从机能否正确丢弃 + 无响应
  2. 超长帧         → 检验接收缓冲区是否溢出
  3. 非法功能码     → 检验从机是否回复 0x01 异常码

用法:
  python3 inject_deformed_packet.py --host 127.0.0.1 --port 5020
  python3 inject_deformed_packet.py --host 127.0.0.1 --port 5020 --mode crc
  python3 inject_deformed_packet.py --host 127.0.0.1 --port 5020 --mode all --count 500
"""

import argparse
import struct
import time
import sys
import random
import socket
import os

# ============================================================
# 工具函数
# ============================================================

def calc_modbus_crc(data: bytes) -> bytes:
    """计算 Modbus CRC16，返回小端序 2 字节"""
    crc = 0xFFFF
    for byte in data:
        crc ^= byte
        for _ in range(8):
            if crc & 0x0001:
                crc = (crc >> 1) ^ 0xA001
            else:
                crc >>= 1
    return struct.pack('<H', crc)


def recv_response(sock: socket.socket, timeout: float = 0.5) -> bytes:
    """接收从机响应，超时返回空 bytes"""
    sock.settimeout(timeout)
    try:
        resp = b''
        while True:
            chunk = sock.recv(256)
            if not chunk:
                break
            resp += chunk
            # Modbus 响应最小 5 字节（异常码）或 7 字节（正常）
            if len(resp) >= 5:
                break
    except socket.timeout:
        pass
    return resp


def format_hex(data: bytes) -> str:
    """bytes 转可读十六进制"""
    return ' '.join(f'{b:02X}' for b in data)


def print_result(index: int, label: str, sent: bytes, resp: bytes, status: str):
    """统一格式打印测试结果（与飞书卡片风格一致）"""
    print(f"┌─ [{index:04d}] {label}")
    print(f"│  发送: {format_hex(sent)}")
    if resp:
        print(f"│  接收: {format_hex(resp)}")
    else:
        print(f"│  接收: (无响应 / 超时)")
    print(f"│  判定: {status}")
    print(f"└────────────────────────────────────────")


# ============================================================
# 畸形包生成器
# ============================================================

def packet_crc_error(slave_addr: int = 0x02) -> bytes:
    """畸形 1: CRC 校验码错误"""
    pdu = bytes([slave_addr, 0x03, 0x00, 0x10, 0x00, 0x01])
    crc = calc_modbus_crc(pdu)
    # 故意取反 CRC 低字节
    corrupted = struct.pack('<H', (~struct.unpack('<H', crc)[0]) & 0xFFFF)
    return pdu + corrupted


def packet_overflow(slave_addr: int = 0x02, extra_len: int = 13) -> bytes:
    """畸形 2: 超长非法数据帧"""
    pdu = bytes([slave_addr, 0x10, 0x00, 0x20, 0x00, 0x02, 0x04])
    # 正常数据是 4 字节，额外填充垃圾数据
    pdu += b'\xDE\xAD\xBE\xEF'  # 正常 4 字节
    pdu += os.urandom(extra_len)
    crc = calc_modbus_crc(pdu)
    return pdu + crc


def packet_illegal_func(slave_addr: int = 0x02) -> bytes:
    """畸形 3: 非法功能码 0x44"""
    pdu = bytes([slave_addr, 0x44, 0x00, 0x10, 0x00, 0x01])
    crc = calc_modbus_crc(pdu)
    return pdu + crc


def packet_noise_burst(count: int = 20) -> bytes:
    """附加: 总线噪声碎片，模拟雷击/逆变器启停干扰"""
    return bytes([random.randint(0x00, 0xFF) for _ in range(count)])


# ============================================================
# 测试运行器
# ============================================================

def run_test(args):
    host = args.host
    port = args.port
    mode = args.mode
    count = args.count
    interval = args.interval

    print("╔══════════════════════════════════════════════════════════╗")
    print("║    Modbus 畸形包注入测试 — STM32 从机鲁棒性验证       ║")
    print(f"║    目标: {host}:{port}                                     ")
    print(f"║    模式: {mode}  次数: {count}  间隔: {interval}s         ")
    print("╚══════════════════════════════════════════════════════════╝")
    print()

    ok_count = 0
    fail_count = 0

    for i in range(1, count + 1):
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(3.0)
            sock.connect((host, port))
        except Exception as e:
            print(f"[{i:04d}] ❌ 连接失败: {e}")
            time.sleep(interval)
            continue

        # ── 选择报文类型 ──
        if mode == 'crc' or (mode == 'all' and i % 3 == 1):
            packet = packet_crc_error()
            label = "CRC错误包"
        elif mode == 'overflow' or (mode == 'all' and i % 3 == 2):
            packet = packet_overflow()
            label = "超长帧注入"
        elif mode == 'illegal' or (mode == 'all' and i % 3 == 0):
            packet = packet_illegal_func()
            label = "非法功能码"
        elif mode == 'noise':
            packet = packet_noise_burst()
            label = "总线噪声碎片"
        else:
            packet = packet_crc_error()
            label = "CRC错误包"

        try:
            sock.sendall(packet)
            resp = recv_response(sock, timeout=args.timeout)

            # ── 结果判定 ──
            if not resp:
                # 无响应 = 从机正确丢弃了错误帧（理想行为）
                status = "✅ 从机正确丢弃，无响应"
                ok_count += 1
            elif len(resp) >= 5 and resp[1] & 0x80:
                # 异常响应码 (MSB of func code = 1)
                exc_code = resp[2]
                if exc_code == 0x01:
                    status = f"✅ 从机返回异常码 0x01 (Illegal Function) — 合规"
                    ok_count += 1
                elif exc_code == 0x02:
                    status = f"✅ 从机返回异常码 0x02 (Illegal Data Address) — 合规"
                    ok_count += 1
                elif exc_code == 0x03:
                    status = f"✅ 从机返回异常码 0x03 (Illegal Data Value) — 合规"
                    ok_count += 1
                else:
                    status = f"⚠️ 从机返回未知异常码 0x{exc_code:02X}"
                    ok_count += 1
            elif len(resp) >= 5:
                # 有响应且无异常 → 从机把畸形包当成有效请求处理了（有风险）
                status = "⚠️  从机响应了畸形包（可能未正确过滤）"
                fail_count += 1
            else:
                status = "❓ 响应过短，无法判定"
                fail_count += 1

            print_result(i, label, packet, resp, status)

        except socket.timeout:
            # 发送后超时 → 从机正确丢弃
            status = "✅ 从机正确丢弃（发送后无响应）"
            print_result(i, label, packet, b'', status)
            ok_count += 1
        except Exception as e:
            status = f"❌ 异常: {e}"
            print_result(i, label, packet, b'', status)
            fail_count += 1
        finally:
            sock.close()

        time.sleep(interval)

    # ── 汇总 ──
    print()
    print("╔══════════════════════════════════════════════════════════╗")
    print("║                      测试报告                           ║")
    print("╠══════════════════════════════════════════════════════════╣")
    print(f"║  总测试次数: {count}                                   ")
    print(f"║  通过 (合规丢弃/异常响应): {ok_count}                   ")
    print(f"║  异常 (错误响应): {fail_count}                          ")
    total_ok_pct = (ok_count / count) * 100
    if ok_count == count:
        print("║  ✅ 结论: 从机鲁棒性通过！所有畸形包被正确丢弃。       ║")
    elif fail_count / count > 0.1:
        print("║  ❌ 结论: 从机存在安全风险，建议检查接收处理逻辑。    ║")
    else:
        print("║  ⚠️  结论: 偶发异常，建议审查特定报文类型。          ║")
    print("╚══════════════════════════════════════════════════════════╝")


# ============================================================
# 命令行入口
# ============================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Modbus 畸形包注入测试 — STM32 从机鲁棒性验证'
    )
    parser.add_argument('--host',  default='127.0.0.1',
                        help='frp 穿透后的目标地址 (默认 127.0.0.1)')
    parser.add_argument('--port',  type=int, default=5020,
                        help='frp 穿透后的目标端口 (默认 5020)')
    parser.add_argument('--mode',  default='all',
                        choices=['crc', 'overflow', 'illegal', 'noise', 'all'],
                        help='测试模式: crc/overflow/illegal/noise/all (默认 all)')
    parser.add_argument('--count',     type=int,   default=100,
                        help='测试发送次数 (默认 100)')
    parser.add_argument('--interval',  type=float, default=0.1,
                        help='每次发送间隔秒数 (默认 0.1)')
    parser.add_argument('--timeout',   type=float, default=0.5,
                        help='响应超时秒数 (默认 0.5)')
    args = parser.parse_args()

    # 修复 packet_overflow 中引用了 os 的问题
    run_test(args)