315 lines
11 KiB
Python
Executable File
315 lines
11 KiB
Python
Executable File
#!/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 # 添加 select 模块导入
|
|
import signal
|
|
|
|
|
|
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._lock = threading.Lock() # Add thread safety
|
|
self.heartbeat_thread: threading.Thread | None = None # Add thread reference
|
|
|
|
def request(self) -> str:
|
|
"""Request license, block until license is obtained or user interrupts"""
|
|
response = self._request_license()
|
|
if response['status'] == 'refuse':
|
|
log.error(response['msg'])
|
|
return None
|
|
elif response['status'] == 'wait_dispatch':
|
|
# Start heartbeat thread
|
|
log.info("Waiting for license allocation...")
|
|
self._start_dispatcher_client(response['dispatcher_server']['port'])
|
|
lic = self.get_lic()
|
|
return lic
|
|
|
|
def get_lic(self) -> str:
|
|
"""Get the assigned license"""
|
|
self.current_lic = self.lic.get()
|
|
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": "refuse", "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)
|
|
self.dispatcher_socket.connect((SERVER_HOST, port))
|
|
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.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)
|
|
print(f"\rwaiting... {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 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, daemon=True)
|
|
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)
|
|
|
|
with self._lock:
|
|
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}")
|
|
|
|
# Global flag for exit control
|
|
is_exiting = False
|
|
clients: dict[str, LicenseDispatcherClient] = {}
|
|
|
|
def signal_handler(signum, frame):
|
|
"""Handle interrupt signals"""
|
|
global is_exiting
|
|
if is_exiting:
|
|
log.info("\nForce exiting...")
|
|
sys.exit(1)
|
|
|
|
log.info("\nReceived interrupt signal, cleaning up...")
|
|
is_exiting = True
|
|
cleanup_and_exit()
|
|
|
|
def cleanup_and_exit():
|
|
"""Cleanup all resources and exit"""
|
|
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")
|
|
sys.exit(0)
|
|
|
|
def main() -> None:
|
|
# Register signal handlers
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
# Make sure SIGINT is not ignored
|
|
if os.name == 'posix':
|
|
signal.siginterrupt(signal.SIGINT, True)
|
|
|
|
while not is_exiting:
|
|
show_menu()
|
|
# Set timeout for input operations
|
|
cmd = input("Please select an operation: ").strip()
|
|
|
|
if is_exiting: # Check if exit flag was set during input
|
|
break
|
|
|
|
if cmd == '1':
|
|
create_client()
|
|
elif cmd == '2':
|
|
release_client()
|
|
elif cmd == '3':
|
|
show_clients()
|
|
elif cmd == '4':
|
|
break
|
|
else:
|
|
print("Invalid command")
|
|
if not is_exiting:
|
|
cleanup_and_exit()
|
|
|
|
def create_client():
|
|
"""Create new license client"""
|
|
global is_exiting
|
|
username = input("Please enter username: ").strip()
|
|
hostname = input("Please enter hostname: ").strip()
|
|
if is_exiting:
|
|
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)
|
|
lic = client.request()
|
|
if lic:
|
|
clients[client_key] = client
|
|
# Add mock license here
|
|
mock_license_add_user(lic, username, hostname)
|
|
print(f"{client_key} successfully created and requested license: {lic}")
|
|
else:
|
|
print(f"{client_key} License request failed")
|
|
|
|
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() |