client.py

import atexit
import ctypes
import pickle
import signal
import socket
import sys
import threading
import tkinter as tk
from tkinter import messagebox, simpledialog

import cv2
import numpy as np
import pyautogui

from protocol import MAX_CONTROL_SIZE, recv_packet, send_packet


DEFAULT_PORT = 9999
CONNECT_TIMEOUT = 10
RECONNECT_DELAY = 3
FRAME_INTERVAL = 0.05
JPEG_ENCODE_PARAMS = [int(cv2.IMWRITE_JPEG_QUALITY), 80]
HD_SIZE = (1280, 720)
FULL_HD_SIZE = (1920, 1080)


def get_transmission_size(screen_width, screen_height):
    """Return the frame size used by the screen-transfer protocol."""
    if screen_width >= FULL_HD_SIZE[0] and screen_height >= FULL_HD_SIZE[1]:
        return FULL_HD_SIZE
    return HD_SIZE


def scale_coordinates(x, y, screen_width, screen_height):
    """Map transmitted-frame coordinates to the local screen."""
    frame_width, frame_height = get_transmission_size(
        screen_width,
        screen_height
    )
    return (
        int(x * screen_width / frame_width),
        int(y * screen_height / frame_height)
    )


class ScreenShareClient:
    def __init__(self):
        self.socket = None
        self.stop_event = threading.Event()
        self.connection_lost = threading.Event()
        self.reconnect_delay = RECONNECT_DELAY
        self.socket_lock = threading.Lock()
        self.server_ip = None
        self.server_port = None
        
        # Регистрируем обработчики для корректного завершения
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        atexit.register(self.cleanup)
        
        # Показываем диалог подключения только один раз
        if len(sys.argv) >= 2:
            self.server_ip = sys.argv[1]

            try:
                # Если порт не указан, используем 9999
                self.server_port = (
                    int(sys.argv[2]) if len(sys.argv) >= 3 else DEFAULT_PORT
                )

                if not 1 <= self.server_port <= 65535:
                    raise ValueError

            except ValueError:
                messagebox.showerror(
                    "Ошибка",
                    "Порт должен быть числом от 1 до 65535"
                )
                sys.exit(1)

        # Аргументы не переданы — показываем диалог
        else:
            if not self.show_connection_dialog():
                sys.exit(0)
            
        # Скрываем консольное окно и приложение после ввода данных
        self.hide_console()
        
    def signal_handler(self, signum, frame):
        """Обработчик сигналов для корректного завершения"""
        print(f"Received signal {signum}, shutting down...")
        self.cleanup()
        sys.exit(0)
        
    def hide_console(self):
        """Скрывает консольное окно"""
        try:
            # Получаем handle консольного окна
            console_window = ctypes.windll.kernel32.GetConsoleWindow()
            if console_window:
                # Скрываем окно
                ctypes.windll.user32.ShowWindow(console_window, 0)
                
                # Убираем окно из панели задач
                ctypes.windll.user32.SetWindowLongW(console_window, -20, 
                                    ctypes.windll.user32.GetWindowLongW(console_window, -20) | 0x00000080)
        except (AttributeError, OSError):
            pass
            
    def show_connection_dialog(self):
        """Показывает диалог для ввода IP и порта"""
        try:
            # Создаем скрытое окно для диалога
            root = tk.Tk()
            root.withdraw()  # Скрываем основное окно
            
            # Показываем диалог
            server_ip = simpledialog.askstring("Connection", 
                                             "Enter server IP address:",
                                             parent=root)
            if not server_ip:
                root.destroy()
                return False
                
            server_port = simpledialog.askstring("Connection", 
                                               "Enter server port (default: 9999):",
                                               parent=root)
            if not server_port:
                server_port = str(DEFAULT_PORT)

            try:
                self.server_port = int(server_port)
                if not 1 <= self.server_port <= 65535:
                    raise ValueError
            except ValueError:
                messagebox.showerror(
                    "Error",
                    "Port must be a number between 1 and 65535"
                )
                root.destroy()
                return False
                
            self.server_ip = server_ip
            root.destroy()
            return True
            
        except (tk.TclError, OSError):
            return False

    def connect_to_server(self):
        """Выполняет одну попытку подключения к серверу."""
        new_socket = None

        try:
            new_socket = socket.socket(
                socket.AF_INET,
                socket.SOCK_STREAM
            )

            new_socket.settimeout(CONNECT_TIMEOUT)
            new_socket.connect(
                (self.server_ip, self.server_port)
            )
            new_socket.settimeout(None)

            new_socket.setsockopt(
                socket.SOL_SOCKET,
                socket.SO_KEEPALIVE,
                1
            )

            with self.socket_lock:
                self.socket = new_socket

            self.connection_lost.clear()

            print(
                f"[CLIENT] Connected to server "
                f"{self.server_ip}:{self.server_port}"
            )

            return True

        except OSError as e:
            if new_socket is not None:
                try:
                    new_socket.close()
                except OSError:
                    pass

            print(
                f"[CLIENT] Connection failed: {e}. "
                f"Retry in {self.reconnect_delay} seconds..."
            )

            return False

    def close_connection(self):
        """Закрывает только текущее соединение, не останавливая клиент."""
        with self.socket_lock:
            current_socket = self.socket
            self.socket = None

        if current_socket is not None:
            try:
                current_socket.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass

            try:
                current_socket.close()
            except OSError:
                pass

    def mark_connection_lost(self, reason=None):
        """Сообщает основному циклу, что требуется переподключение."""
        if self.connection_lost.is_set():
            return

        if reason and not self.stop_event.is_set():
            print(f"[CLIENT] Connection lost: {reason}")

        self.connection_lost.set()
        self.close_connection()

    def send_screen(self):
        """Передаёт экран до обрыва текущего соединения."""
        try:
            while (
                    not self.stop_event.is_set()
                    and not self.connection_lost.is_set()
            ):
                with self.socket_lock:
                    current_socket = self.socket

                if current_socket is None:
                    break

                frame = np.asarray(pyautogui.screenshot())
                frame = cv2.cvtColor(
                    frame,
                    cv2.COLOR_RGB2BGR
                )

                screen_height, screen_width = frame.shape[:2]
                target_size = get_transmission_size(
                    screen_width,
                    screen_height
                )

                if (screen_width, screen_height) != target_size:
                    is_downscaling = (
                        target_size[0] < screen_width
                        or target_size[1] < screen_height
                    )
                    interpolation = (
                        cv2.INTER_AREA if is_downscaling else cv2.INTER_LINEAR
                    )
                    frame = cv2.resize(
                        frame,
                        target_size,
                        interpolation=interpolation
                    )

                success, buffer = cv2.imencode(
                    ".jpg",
                    frame,
                    JPEG_ENCODE_PARAMS
                )

                if not success:
                    continue

                send_packet(current_socket, memoryview(buffer))

                if self.stop_event.wait(FRAME_INTERVAL):
                    break

        except OSError as e:
            self.mark_connection_lost(e)

        except Exception as e:
            self.mark_connection_lost(
                f"screen transmission error: {e}"
            )

    def receive_control(self):
        """Принимает команды до обрыва текущего соединения."""
        try:
            while (
                    not self.stop_event.is_set()
                    and not self.connection_lost.is_set()
            ):
                with self.socket_lock:
                    current_socket = self.socket

                if current_socket is None:
                    break

                command_data = recv_packet(
                    current_socket,
                    MAX_CONTROL_SIZE
                )

                if not command_data:
                    self.mark_connection_lost(
                        "server closed the connection"
                    )
                    break

                try:
                    command = pickle.loads(command_data)
                    self.execute_command(command)

                except Exception as e:
                    print(
                        f"[CLIENT] Invalid control command: {e}"
                    )

        except (OSError, ValueError) as e:
            self.mark_connection_lost(e)

    def execute_command(self, command):
        """Выполняет команду управления"""
        try:
            cmd_type = command.get('type')
            
            if cmd_type in {
                    'mouse_move',
                    'mouse_click',
                    'mouse_double_click',
                    'mouse_scroll'
            }:
                screen_width, screen_height = pyautogui.size()
                real_x, real_y = scale_coordinates(
                    command['x'],
                    command['y'],
                    screen_width,
                    screen_height
                )

                if cmd_type == 'mouse_move':
                    pyautogui.moveTo(real_x, real_y)
                elif cmd_type == 'mouse_click':
                    pyautogui.click(
                        real_x,
                        real_y,
                        button=command.get('button', 'left')
                    )
                elif cmd_type == 'mouse_double_click':
                    pyautogui.doubleClick(
                        real_x,
                        real_y,
                        button=command.get('button', 'left')
                    )
                else:
                    pyautogui.scroll(
                        command.get('clicks', 1),
                        x=real_x,
                        y=real_y
                    )
                
            elif cmd_type == 'key_press':
                key = command['key']
                pyautogui.press(key)
                
        except Exception as e:
            print("[CLIENT] Error in execute_command:", e)

    def run(self):
        """Подключается и автоматически восстанавливает соединение."""
        print(
            f"[CLIENT] Target server: "
            f"{self.server_ip}:{self.server_port}"
        )

        try:
            while not self.stop_event.is_set():
                if not self.connect_to_server():
                    if self.stop_event.wait(self.reconnect_delay):
                        break
                    continue

                print("[CLIENT] Starting screen and control threads...")

                screen_thread = threading.Thread(
                    target=self.send_screen,
                    daemon=True
                )

                control_thread = threading.Thread(
                    target=self.receive_control,
                    daemon=True
                )

                screen_thread.start()
                control_thread.start()

                while not self.stop_event.is_set():
                    if self.connection_lost.wait(0.2):
                        break

                self.close_connection()

                screen_thread.join(timeout=2)
                control_thread.join(timeout=2)

                if not self.stop_event.is_set():
                    print(
                        f"[CLIENT] Reconnecting in "
                        f"{self.reconnect_delay} seconds..."
                    )

                    if self.stop_event.wait(self.reconnect_delay):
                        break

        except KeyboardInterrupt:
            print("[CLIENT] Interrupted by user.")

        finally:
            self.cleanup()

    def cleanup(self):
        """Окончательно останавливает клиент."""
        self.stop_event.set()
        self.connection_lost.set()
        self.close_connection()

if __name__ == "__main__":
    client = ScreenShareClient()
    client.run()