from gmssl import sm4
from cryptography.hazmat.primitives.asymmetric import dh
from cryptography.hazmat.primitives import serialization
import socket
import sys
import threading
from controller import main

PASSWORD = "afaksodfjasdfcoi"

RFC3526_PRIME_2048 = 0xFFFFFFFFFFFFFFFFC90FDAA22168C234C4C6628B80DC1CD129024E088A67CC74020BBEA63B139B22514A08798E3404DDEF9519B3CD3A431B302B0A6DF25F14374FE1356D6D51C245E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7EDEE386BFB5A899FA5AE9F24117C4B1FE649286651ECE45B3DC2007CB8A163BF0598DA48361C55D39A69163FA8FD24CF5F83655D23DCA3AD961C62F356208552BB9ED529077096966D670C354E4ABC9804F1746C08CA18217C32905E462E36CE3BE39E772C180E86039B2783A2EC07A28FB5C55DF06F4C52C9DE2BCBF6955817183995497CEA956AE515D2261898FA051015728E5A8AACAA68FFFFFFFFFFFFFFFF
RFC3526_GENERATOR = 2

def generate_dh_parameters():
    parameter_numbers = dh.DHParameterNumbers(
        p=RFC3526_PRIME_2048,
        g=RFC3526_GENERATOR
    )
    return parameter_numbers.parameters()

def server_dh_key_agreement(conn, parameters):
    """服务端DH密钥交换"""
    # 生成服务端DH密钥对
    server_private_key = parameters.generate_private_key()
    
    # 发送公钥
    server_public_key = server_private_key.public_key()
    server_public_bytes = server_public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    conn.sendall(server_public_bytes)
    
    # 接收客户端公钥
    client_public_bytes = conn.recv(1024)
    client_public_key = serialization.load_pem_public_key(
        client_public_bytes,
    )
    
    # 计算共享密钥
    shared_key = server_private_key.exchange(client_public_key)
    return shared_key[:16]  # SM4需要16字节密钥

def sm4_encrypt(data, key):
    """SM4加密"""
    crypt_sm4 = sm4.CryptSM4()
    crypt_sm4.set_key(key, sm4.SM4_ENCRYPT)
    return crypt_sm4.crypt_ecb(data)

def sm4_decrypt(data, key):
    """SM4解密"""
    crypt_sm4 = sm4.CryptSM4()
    crypt_sm4.set_key(key, sm4.SM4_DECRYPT)
    return crypt_sm4.crypt_ecb(data)

# class SecureSocketIO:
#     """安全的Socket I/O包装器"""
#     def __init__(self, conn, key):
#         self.conn = conn
#         self.key = key
#         self.buffer = b""
    
#     def write(self, data):
#         """加密并发送数据"""
#         try:
#             encrypted = sm4_encrypt(data.encode('utf-8'), self.key)
#             self.conn.sendall(len(encrypted).to_bytes(4, 'big') + encrypted)
#         except Exception as e:
#             print(f"加密发送错误: {e}", file=sys.__stdout__)
    
#     def readline(self):
#         """接收并解密数据"""
#         self.write('__!input()_called!__')
#         while True:
#             # 检查缓冲区是否有完整消息
#             if len(self.buffer) >= 4:
#                 msg_len = int.from_bytes(self.buffer[:4], 'big')
#                 if len(self.buffer) >= 4 + msg_len:
#                     encrypted = self.buffer[4:4+msg_len]
#                     self.buffer = self.buffer[4+msg_len:]
#                     try:
#                         decrypted = sm4_decrypt(encrypted, self.key)
#                         return decrypted.decode('utf-8') + '\n'
#                     except Exception as e:
#                         print(f"解密错误: {e}", file=sys.__stdout__)
#                         return '\n'
            
#             # 接收更多数据
#             try:
#                 data = self.conn.recv(4096)
#                 if not data:
#                     return ''
#                 self.buffer += data
#             except Exception as e:
#                 print(f"接收错误: {e}", file=sys.__stdout__)
#                 return ''
    
#     def flush(self):
#         pass

class SecureSocketWriter:
    """线程安全的加密socket写入器"""
    def __init__(self, conn, key):
        self.conn = conn
        self.key = key
        self.lock = threading.Lock()  # 写入锁保证线程安全
    
    def write(self, data):
        try:
            with self.lock:
                encrypted = sm4_encrypt(data.encode('utf-8'), self.key)
                # 添加长度前缀
                self.conn.sendall(len(encrypted).to_bytes(4, 'big') + encrypted)
        except (BrokenPipeError, ConnectionResetError):
            # 连接已断开
            pass
        except Exception as e:
            print(f"加密写入错误: {e}", file=sys.__stderr__)
    
    def flush(self):
        pass

class SecureSocketReader:
    """线程安全的解密socket读取器"""
    def __init__(self, conn, key):
        self.conn = conn
        self.key = key
        self.buffer = bytearray()
        self.lock = threading.Lock()  # 读取锁保证线程安全
    
    def readline(self):
        with self.lock:
            print('__!input()_called!__')
            while True:
                # 检查缓冲区是否有完整消息
                if len(self.buffer) >= 4:
                    msg_len = int.from_bytes(self.buffer[:4], 'big')
                    if len(self.buffer) >= 4 + msg_len:
                        encrypted = self.buffer[4:4+msg_len]
                        self.buffer = self.buffer[4+msg_len:]
                        try:
                            decrypted = sm4_decrypt(encrypted, self.key)
                            return decrypted.decode('utf-8') + '\n'
                        except Exception as e:
                            print(f"解密错误: {e}", file=sys.__stdout__)
                            return '\n'
                
                # 接收新数据
                try:
                    data = self.conn.recv(4096)
                    if not data:
                        return ''
                    self.buffer += data
                except Exception as e:
                    print(f"接收错误: {e}", file=sys.__stdout__)
                    return ''

class ThreadLocalIO:
    """管理线程本地I/O重定向"""
    def __init__(self):
        self.local = threading.local()
    
    def register(self, reader, writer):
        self.local.reader = reader
        self.local.writer = writer
    
    def get_reader(self):
        return getattr(self.local, 'reader', sys.__stdin__)
    
    def get_writer(self):
        return getattr(self.local, 'writer', sys.__stdout__)

# 全局I/O管理器
thread_io_manager = ThreadLocalIO()

class ThreadSafePrint:
    def write(self, text):
        writer = thread_io_manager.get_writer()
        if hasattr(writer, 'write'):
            writer.write(text)
        else:
            sys.__stdout__.write(text)
    
    def flush(self):
        writer = thread_io_manager.get_writer()
        if hasattr(writer, 'flush'):
            writer.flush()
        else:
            sys.__stdout__.flush()

class ThreadSafeInput:
    def readline(self):
        reader = thread_io_manager.get_reader()
        if hasattr(reader, 'readline'):
            return reader.readline()
        return sys.__stdin__.readline()

# 重定向全局I/O
sys.stdout = ThreadSafePrint()
sys.stderr = ThreadSafePrint()
sys.stdin = ThreadSafeInput()

class ClientHandler(threading.Thread):
    """客户端处理线程"""
    def __init__(self, conn, addr, key):
        threading.Thread.__init__(self)
        self.conn = conn
        self.addr = addr
        self.key = key
        self.reader = SecureSocketReader(conn, key)
        self.writer = SecureSocketWriter(conn, key)
        self.daemon = True  # 主线程退出时自动结束
    
    def run(self):
        try:
            # 注册线程本地I/O
            thread_io_manager.register(self.reader, self.writer)
            
            # 认证流程
            sys.stdout.write("请输入密码: ")
            password = sys.stdin.readline().strip()
            
            if password == PASSWORD:
                sys.stdout.write("认证成功!\n")
                # 调用控制器主函数
                main(True)
            else:
                sys.stdout.write("认证失败，连接关闭\n")
            
        except Exception as e:
            print(f"处理客户端 {self.addr} 时出错: {e}", file=sys.__stdout__)
        finally:
            self.conn.close()
            print(f"与 {self.addr} 的连接已关闭", file=sys.__stdout__)

def start_server(host='0.0.0.0', port=65432):
    """启动服务器"""
    parameters = generate_dh_parameters()
    
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((host, port))
        s.listen()
        s.settimeout(1)
        print(f"服务器启动，监听 {host}:{port}")
        
        try:
            while True:
                try:
                    conn, addr = s.accept()
                    print(f"接收到来自 {addr} 的新连接", file=sys.__stdout__)

                    # 执行DH密钥交换
                    shared_key = server_dh_key_agreement(conn, parameters)
                    print(f"与 {addr} 建立共享密钥: {shared_key.hex()}", file=sys.__stdout__)
                    
                    # 为每个客户端创建处理线程
                    handler = ClientHandler(conn, addr, shared_key)
                    handler.start()
                except socket.timeout:
                    continue
                
        except KeyboardInterrupt:
            print("正在关闭服务器...")
        finally:
            s.close()

if __name__ == "__main__":
    start_server()