vcast_license_usage_server/license_client.py

#!/usr/bin/python3
import json
import os
import socket
import threading
import time
from datetime import datetime
import http.client
import sys
from my_log import get_logger
from mock_lmutil import mock_license_add_user, mock_license_remove_user, _read_sample_file, _write_sample_file
import queue
import select


log = get_logger()

server = os.environ.get('SSTAR_VCAST_LICENSE_SERVER_HOST')  # Use get() to avoid KeyError
SERVER_HOST = server if server else '127.0.0.1'
SERVER_PORT = 8088
SERVER_URL = f'http://{SERVER_HOST}:{SERVER_PORT}'

class LicenseDispatcherClient:
    def __init__(self, user: str = '', host: str = '') -> None:
        self.server_url = SERVER_URL  # Use global configuration
        self.user = user
        self.host = host
        self.dispatcher_socket: socket.socket = None
        self.exit_event = threading.Event()
        self.lic: queue.Queue = queue.Queue()
        self.current_lic: str = None
        self.dispatcher_timer: threading.Timer = None  # Add missing attribute
        self.heartbeat_thread: threading.Thread | None = None  # Add thread reference

    def request(self) -> str:
        """Request license, block until license is obtained"""
        response = self._request_license()
        if response['status'] == 'wait_dispatch':
            # Start heartbeat thread
            log.info("Waiting for license allocation...")
            self._start_dispatcher_client(response['dispatcher_server']['port'])
        else:
            log.warning(f"{response['status']}: {response['msg']}")
        return response['status']

    def get_lic(self) -> str:
        """Get the assigned license"""
        while not self.exit_event.is_set():
            try:
                self.current_lic = self.lic.get(timeout=0.1)
                return self.current_lic
            except queue.Empty:
                continue
        return self.current_lic

    def _request_license(self) -> dict:
        """Send license request to server"""
        conn = None
        try:
            log.info(f"Requesting license: {self.user}@{self.host}")
            conn = http.client.HTTPConnection(SERVER_HOST, SERVER_PORT, timeout=10)  # Add timeout
            headers = {'Content-Type': 'application/json'}
            data = {
                "user": self.user,
                "host": self.host
            }
            conn.request('GET', '/vcast/request_lic',
                        body=json.dumps(data),
                        headers=headers)
            response = conn.getresponse()
            response_data = response.read().decode()
            return json.loads(response_data)
        except (socket.timeout, ConnectionRefusedError) as e:
            log.error(f"Connection error: {e}")
            return {"status": "connect_error", "msg": str(e)}
        except Exception as e:
            log.error(f"Requesting license error: {e}")
            return {"status": "error", "msg": str(e)}
        finally:
            if conn:
                conn.close()

    def _start_dispatcher_client(self, port: int) -> None:
        """Start dispatcher connection with server"""
        log.info(f"Starting dispatcher: {port}")
        self.dispatcher_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.dispatcher_socket.connect((SERVER_HOST, port))
        except Exception as e:
            log.error(f"Failed to connect to dispatcher: {e}")
            return
        start_time = time.time()

        def heartbeat_loop():
            while not self.exit_event.is_set():
                try:
                    # Send heartbeat
                    heartbeat = {
                        "type": "heartbeat",
                        "user": self.user,
                        "host": self.host,
                        "heartbeat_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    }
                    self.dispatcher_socket.send(json.dumps(heartbeat).encode('utf-8'))

                    # Wait for response with timeout
                    readable, _, _ = select.select([self.dispatcher_socket], [], [], 0.5)
                    if self.exit_event.is_set():
                        break
                    if self.dispatcher_socket in readable:
                        data = self.dispatcher_socket.recv(1024)
                        if not data:
                            log.warning("Connection closed by server")
                            break
                        response = json.loads(data.decode())
                        if response.get('type') == 'lic_dispatch':
                            self.lic.put(response['lic'])
                            log.info(f"Received license: {response['lic']}")
                            # Send acknowledgment
                            ack = {
                                "type": "lic_received_ack",
                                "user": self.user,
                                "host": self.host
                            }
                            self.dispatcher_socket.send(json.dumps(ack).encode('utf-8'))
                            break
                        elif response.get('type') == 'heartbeat_ack':
                            wait_time = int(time.time() - start_time)
                            if wait_time < 60:
                                print(f"\rwaiting... {wait_time}s", end='', flush=True)
                            elif wait_time < 3600:
                                print(f"\rwaiting... {wait_time / 60}m {wait_time}s", end='', flush=True)
                            else:
                                print(f"\rwaiting... {wait_time / 3600}h {wait_time / 60}m {wait_time}s", end='', flush=True)
                        elif response.get('type') == 'queue':
                            wait_time = int(time.time() - start_time)
                            log.info(f"waiting... {wait_time}s")

                    time.sleep(0.5)  # Small sleep to prevent CPU spinning
                except KeyboardInterrupt:
                    log.info("Keyboard interrupt received, shutting down...")
                    break
                except Exception as e:
                    log.error(f"Error in heartbeat loop: {e}")
                    break

            self.lic.put(None)  # Signal end of heartbeat loop

        # Start heartbeat thread
        self.heartbeat_thread = threading.Thread(target=heartbeat_loop)
        self.heartbeat_thread.start()

    def release(self) -> None:
        """Release the current license"""
        log.info("License released")

    def exit(self) -> None:
        """Exit the client and cleanup resources"""
        try:
            self.exit_event.set()
            
            # Wait for heartbeat thread to finish with timeout
            if self.heartbeat_thread and self.heartbeat_thread.is_alive():
                self.heartbeat_thread.join(timeout=1.0)
            
            if self.dispatcher_socket:
                try:
                    self.dispatcher_socket.shutdown(socket.SHUT_RDWR)
                except Exception:
                    pass
                self.dispatcher_socket.close()
                self.dispatcher_socket = None
            
            self.release()
            self.lic.put(None)
        except Exception as e:
            log.error(f"Error during exit: {e}")

clients: dict[str, LicenseDispatcherClient] = {}
def cleanup_clients():
    """清理所有客户端资源"""
    for client_key, client in list(clients.items()):
        try:
            log.info(f"Cleaning up client: {client_key}")
            client.exit()
            del clients[client_key]
        except Exception as e:
            log.error(f"Error cleaning up client {client_key}: {e}")
    log.info("\nProgram exited")

def main() -> None:
    try:
        while True:
            show_menu()
            cmd = input("Please select an operation: ").strip()

            if cmd == '1':
                create_client()
            elif cmd == '2':
                release_client()
            elif cmd == '3':
                show_clients()
            elif cmd == '4':
                break
            else:
                print("Invalid command")
    except KeyboardInterrupt:
        print("\nReceived keyboard interrupt, cleaning up...")
    finally:
        cleanup_clients()

def create_client():
    """Create new license client"""
    try:
        username = input("Please enter username: ").strip()
        hostname = input("Please enter hostname: ").strip()
    except KeyboardInterrupt:
        print("\nOperation cancelled")
        return

    if not username or not hostname:
        print("Username and hostname cannot be empty")
        return

    client_key = f"{username}@{hostname}"
    if client_key in clients:
        print(f"Client {client_key} already exists")
        return

    client = LicenseDispatcherClient(user=username, host=hostname)
    status = client.request()
    if status == 'wait_dispatch':
        lic = client.get_lic()
        clients[client_key] = client
        mock_license_add_user(lic, username, hostname)
        print(f"{client_key} successfully created and requested license: {lic}")
    else:
        if status == 'connect_error':
            print(f"{client_key} License Server connect error, You can try again later")
        else:
            print(f"{client_key} License request failed: {status}")

def release_client():
    """Release existing license"""
    lic_use_info = _read_sample_file()
    if not lic_use_info:
        print("No license usage record currently")
        return

    print("\nCurrent license usage:")
    users_list = []

    # Iterate through all license server information
    for server_info in lic_use_info.values():
        for used_info in server_info.get("used_info", []):
            user = used_info["user"]
            host = used_info["host"]
            lic = used_info["lic"]
            users_list.append((user, host, lic))
            print(f"{len(users_list)}: user: {user}, host: {host}, lic: {lic}")

    if not users_list:
        print("No valid user records found")
        return

    try:
        choice = int(input("\nPlease select the license number to release (0 to cancel): "))
    except ValueError:
        print("Please enter a valid number")
    if choice == 0:
        return
    if 1 <= choice <= len(users_list):
        user, host, lic = users_list[choice-1]
        client_key = f"{user}@{host}"
        lic_use_info = _read_sample_file()

        # If the client exists in the clients dictionary, release it
        if client_key in clients:
            client = clients[client_key]
            client.release()
            client.exit()
            del clients[client_key]

        # Remove mock license record
        mock_license_remove_user(user, host)
        print(f"Released license: user: {user}, host: {host}, lic: {lic}")
    else:
        print("Invalid selection")

def show_clients():
    if not clients:
        print("No active clients currently")
        return

    print("\nCurrent active clients:")
    for key, client in clients.items():
        lic_status = "License obtained" if client.current_lic else "Waiting"
        print(f"- {key}: {lic_status}")

def show_menu() -> None:
    print("\n=== License Management System ===")
    print("1: Create new license request")
    print("2: Release license")
    print("3: Show current clients")
    print("4: Exit")

if __name__ == "__main__":
    main()