from scapy.all import sniff, IP, UDP
import socket
import subprocess

# 获取本机IP
def get_local_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
        return local_ip
    except:
        return "127.0.0.1"

TARGET_PORT = 4705
MY_IP = get_local_ip()
blocked_ips = set()  # 已拉黑IP集合，避免重复操作

# 冷麟攻击特征码
LENG_LIN_MAGIC_1 = bytes.fromhex("444d4f4300000100")
LENG_LIN_MAGIC_2 = bytes.fromhex("63006d0064002e00650078006500")
LENG_LIN_MAGIC_3 = bytes.fromhex("2f006b002000")

def is_lenglin_attack(payload):
    if not payload.startswith(LENG_LIN_MAGIC_1):
        return False
    if LENG_LIN_MAGIC_2 in payload or LENG_LIN_MAGIC_3 in payload:
        return True
    if len(payload) > 200 and b"\x00\x00\x00\x01\x00\x00\x00" in payload:
        return True
    return False

def block_ip(ip):
    """Windows防火墙自动拉黑IP"""
    if ip not in blocked_ips:
        cmd = f'netsh advfirewall firewall add rule name="Block LengLin {ip}" dir=in action=block remoteip={ip}'
        subprocess.run(cmd, shell=True, capture_output=True)
        blocked_ips.add(ip)
        print(f"[✅] 已自动拦截冷麟攻击IP: {ip}")

def detect_attack(packet):
    if IP in packet and UDP in packet:
        if packet[UDP].dport == TARGET_PORT and packet[IP].dst == MY_IP:
            src_ip = packet[IP].src
            payload = packet[UDP].payload.load
            if is_lenglin_attack(payload):
                print(f"[!] 【冷麟恶意攻击】来源IP: {src_ip}")
                block_ip(src_ip)

print("=== 冷麟攻击检测+自动拦截已启动 ===")
print("✅ 教师端正常指令已过滤，仅检测并拦截冷麟攻击\n")
sniff(prn=detect_attack, store=0)