import socket
import threading
import json
import hashlib
import time
import struct


# 定义 Provider 类，用于处理与协调器的连接和P2P通信
class Provider:
    def __init__(self, coordinator_host='127.0.0.1', coordinator_port=5000):
        # 初始化协调器的主机和端口
        self.coordinator_host = coordinator_host
        self.coordinator_port = coordinator_port
        # 用于存储认证令牌
        self.token = None
        # 定义可提供的服务及其默认端口
        self.service_ports = {'ssh': 22, 'alist': 5244, 'minecraft': 25565}
        # 存储连接的客户端
        self.connections = {}
        # 用于线程安全操作的锁
        self.lock = threading.Lock()

    def connect_to_coordinator(self):
        # 创建与协调器的TCP连接
        self.coord_conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.coord_conn.connect((self.coordinator_host, self.coordinator_port))

        # 发送登录请求
        self._send_json({'action': 'login', 'account': 'admin'})
        response = self._recv_json()

        # 处理协调器返回的盐值并进行密码哈希验证
        if response.get('status') == 'salt':
            salt = response['salt']
            password_hash = hashlib.sha256((salt + "admin_password").encode()).hexdigest()
            self._send_json({'action': 'auth', 'hash': password_hash})
            response = self._recv_json()

            # 如果认证成功，存储令牌并注册服务
            if response.get('status') == 'success':
                self.token = response['token']
                print(f"Authenticated. Token: {self.token}")

                self._send_json({
                    'action': 'register_service',
                    'services': list(self.service_ports.keys()),
                    'token': self.token
                })
                response = self._recv_json()
                if response.get('status') == 'success':
                    print("Services registered")
                    return True
        print("Connection to coordinator failed")
        return False

    def handle_punch_request(self, data):
        connector_addr = tuple(data['connector_addr'])
        service_name = data['service_name']
        print(f"Punching hole to connector at {connector_addr}, waiting 10 seconds...")

        # Wait for 10 seconds to allow the connector to initiate its punch
        time.sleep(2)

        # 使用UDP打洞
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # 绑定到相同的本地端口（用于后续TCP连接）
        udp_socket.bind(('0.0.0.0', 0))
        punch_port = udp_socket.getsockname()[1]
        # 向对方发送打洞包
        for i in range(10):
            udp_socket.sendto(b'punch', connector_addr)
            time.sleep(0.2)


        punch_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        punch_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        punch_sock.settimeout(10)
        punch_sock.bind(('0.0.0.0', punch_port))

        try:
            punch_sock.connect(connector_addr)
            print("Successfully connected to connector after delay")
            threading.Thread(
                target=self.handle_connector_connection,
                args=(punch_sock, service_name),
                daemon=True
            ).start()
        except socket.error as e:
            print(f"Punching failed: {e}")
            punch_sock.close()

    def handle_connector_connection(self, sock, service_name):
        # 处理与客户端的连接，启动心跳机制
        threading.Thread(target=self.send_heartbeats, args=(sock,), daemon=True).start()

        try:
            while True:
                # 接收连接头信息
                header = sock.recv(5)
                if not header:
                    break

                conn_id, data_len = struct.unpack("!I B", header)
                data = sock.recv(data_len) if data_len > 0 else b''

                if not data:
                    with self.lock:
                        if conn_id in self.connections:
                            self.connections[conn_id].close()
                            del self.connections[conn_id]
                    continue

                with self.lock:
                    if conn_id not in self.connections:
                        service_port = self.service_ports.get(service_name, 22)
                        service_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                        service_sock.connect(('127.0.0.1', service_port))
                        self.connections[conn_id] = service_sock
                        threading.Thread(
                            target=self.forward_data,
                            args=(service_sock, sock, conn_id),
                            daemon=True
                        ).start()

                    self.connections[conn_id].sendall(data)
        except ConnectionResetError:
            pass
        finally:
            sock.close()
            with self.lock:
                for conn_id, service_sock in list(self.connections.items()):
                    service_sock.close()
                self.connections.clear()

    def forward_data(self, src, dst, conn_id):
        # 转发数据
        try:
            while True:
                data = src.recv(4096)
                if not data:
                    break
                header = struct.pack("!I B", conn_id, len(data))
                dst.sendall(header + data)
        finally:
            src.close()
            with self.lock:
                if conn_id in self.connections:
                    del self.connections[conn_id]

    def send_heartbeats(self, sock):
        # 发送心跳包以保持连接
        while True:
            try:
                sock.sendall(b'\x00\x00\x00\x00\x00')  # Empty heartbeat
                time.sleep(5)
            except:
                break

    def start(self):
        # 启动提供者，连接到协调器并开始处理请求
        if not self.connect_to_coordinator():
            return

        try:
            while True:
                data = self._recv_json()
                if data and data.get('action') == 'punch_request':
                    self.handle_punch_request(data)
        except (ConnectionResetError, json.JSONDecodeError):
            print("Disconnected from coordinator")

    def _send_json(self, data):
        # 发送JSON数据
        self.coord_conn.sendall(json.dumps(data).encode())

    def _recv_json(self):
        # 接收JSON数据
        data = self.coord_conn.recv(4096)
        return json.loads(data.decode()) if data else None


if __name__ == "__main__":
    provider = Provider(coordinator_host='www.awin-x.top')
    provider.start()