import zlib
import time
import socket
import struct
from loguru import logger
import mss
import mss.tools
import pyautogui
import configparser
from Crypto.Cipher import AES
import base64
import ssl

class RemoteClient:
    def __init__(self, host='127.0.0.1', port=5000):
        config = configparser.ConfigParser()
        config.read('config.ini')
        
        self.host = config['Server']['host']
        self.port = int(config['Server']['port'])
        self.client_id = config['Authentication']['client_id']
        
        # 初始化加密组件
        self.cipher = AES.new(
            base64.b64decode(config['Security']['encryption_key']), 
            AES.MODE_CBC, 
            base64.b64decode(config['Security']['encryption_salt'])
        )
        
        # 创建SSL上下文
        self.context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        self.context.load_cert_chain(config['Authentication']['cert_file'])
        self.context.check_hostname = False
        
        self.sock = self.context.wrap_socket(
            socket.socket(socket.AF_INET, socket.SOCK_STREAM),
            server_hostname=config['Server']['host']
        )
        self.prev_hash = None
        self.prev_mouse = (0, 0)
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        logger.add('remote_client.log', 
                   rotation='10 MB',
                   retention='7 days',
                   format='{time:YYYY-MM-DD HH:mm:ss} | {level} | {module}:{function} - {message}')

    def connect(self):
        try:
            self.sock.connect((self.host, self.port))
            logger.info(f'成功连接到服务器 {self.host}:{self.port}')
        except Exception as e:
            logger.error(f'连接错误: {e}')

    def capture_screen(self):
        with mss.mss() as sct:
            screenshot = sct.grab(sct.monitors[0])
            raw_bytes = mss.tools.to_png(sct.grab(sct.monitors[0]), (0, 0, screenshot.width, screenshot.height))
        current_hash = zlib.crc32(raw_bytes)
        
        if current_hash != self.prev_hash:
            original_size = len(raw_bytes)
            compressed = zlib.compress(raw_bytes, level=9)
            compressed_size = len(compressed)
            ratio = (1 - compressed_size/original_size) * 100
            logger.debug(f'屏幕变化检测 [尺寸: ({screenshot.width},{screenshot.height})] [压缩率: {ratio:.1f}%]')
            self.prev_hash = current_hash
            return ('screen', compressed, (screenshot.width, screenshot.height))
        return None

    def get_mouse_data(self):
        x, y = pyautogui.position()
        buttons = pyautogui.mouseDown()
        
        if (x, y) != self.prev_mouse or buttons:
            logger.info(f'鼠标事件 [坐标: ({x},{y})] [按键: {buttons}]')
            data = struct.pack('!II?', x, y, buttons)
            self.prev_mouse = (x, y)
            return ('mouse', data)
        return None

    def send_data(self, data_type, payload, size=None):
        header = struct.pack('!4sI', data_type.encode(), len(payload))
        if size:
            header += struct.pack('!II', *size)
        try:
            self.sock.sendall(header + payload)
            logger.trace(f'发送{data_type}数据 [大小: {len(payload)}字节]')
        except Exception as e:
            logger.warning(f'发送错误: {e}')
            self.connect()

    def run(self):
        self.connect()
        while True:
            start_time = time.time()
            
            # 捕获屏幕差异
            screen_data = self.capture_screen()
            if screen_data:
                self.send_data(*screen_data)
            
            # 捕获鼠标状态
            mouse_data = self.get_mouse_data()
            if mouse_data:
                self.send_data(*mouse_data)
            
            # 保持2秒间隔
            elapsed = time.time() - start_time
            if elapsed < 2:
                time.sleep(2 - elapsed)

if __name__ == '__main__':
    client = RemoteClient()
    client.run()