guanghulab/exe-engine/gatekeeper/gatekeeper-client.py

#!/usr/bin/env python3
"""
gatekeeper-client — 成员服务器接入铸渊算力调度池

用法：
  python3 gatekeeper-client.py --register --name YOUR_NAME --master 43.139.217.141:3910
  python3 gatekeeper-client.py --status
"""
import json, os, sys, time, socket, uuid, hashlib, subprocess
from datetime import datetime

CONFIG_DIR = os.path.expanduser('~/.gatekeeper')
CONFIG_FILE = os.path.join(CONFIG_DIR, 'config.json')

def log(m):
    print(f'[{datetime.now().strftime("%H:%M:%S")}] {m}', flush=True)

def ensure_dir():
    os.makedirs(CONFIG_DIR, exist_ok=True)

def load_config():
    try:
        with open(CONFIG_FILE) as f:
            return json.load(f)
    except:
        return {}

def save_config(cfg):
    ensure_dir()
    with open(CONFIG_FILE, 'w') as f:
        json.dump(cfg, f, indent=2)

def register(name, master):
    """向主控闸门注册"""
    log(f'注册到 {master}')
    host, port = master.split(':')
    port = int(port)
    
    # 生成本机身份
    node_id = str(uuid.uuid4())[:8]
    key = hashlib.sha256((node_id + name + str(time.time())).encode()).hexdigest()[:32]
    
    # 通过 socket 发送注册请求
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(10)
        s.connect((host, port))
        req = json.dumps({
            'action': 'register',
            'name': name,
            'node_id': node_id,
            'key': key
        })
        s.sendall((req + '\n').encode())
        resp = s.recv(4096).decode()
        s.close()
        
        data = json.loads(resp)
        if data.get('status') == 'ok':
            save_config({'name': name, 'node_id': node_id, 'key': key, 'master': master, 'status': 'registered'})
            log('注册成功！身份密钥已保存在本地')
            log(f'名称: {name}')
            log(f'节点ID: {node_id}')
            return True
        else:
            log(f'注册失败: {data.get("message", "?")}')
            return False
    except Exception as e:
        log(f'连接失败: {e}')
        log('请检查服务器是否能访问 ' + master)
        return False

def show_status():
    """显示当前连接状态"""
    cfg = load_config()
    if not cfg:
        log('未注册。请先运行: python3 gatekeeper-client.py --register --name YOUR_NAME --master IP:PORT')
        return
    log(f'名称: {cfg.get("name", "?")}')
    log(f'节点ID: {cfg.get("node_id", "?")}')
    log(f'主控闸门: {cfg.get("master", "?")}')
    log(f'状态: {cfg.get("status", "?")}')
    
    # 测试连接
    master = cfg.get('master', '')
    if master:
        try:
            host, port = master.split(':')
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(5)
            s.connect((host, int(port)))
            s.sendall(json.dumps({'action': 'ping', 'node_id': cfg.get('node_id')})['\n'].encode())
            resp = s.recv(4096).decode()
            s.close()
            data = json.loads(resp)
            if data.get('status') == 'ok':
                log('连接: ✅ connected')
            else:
                log('连接: ❌ connection error')
        except Exception as e:
            log(f'连接: ❌ {e}')

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='gatekeeper-client: 接入铸渊算力调度池')
    parser.add_argument('--register', action='store_true', help='注册到主控闸门')
    parser.add_argument('--name', help='你的名称 (如 awen)')
    parser.add_argument('--master', default='43.139.217.141:3910', help='主控闸门地址')
    parser.add_argument('--status', action='store_true', help='查看连接状态')
    args = parser.parse_args()
    
    if args.status:
        show_status()
    elif args.register:
        if not args.name:
            print('请提供 --name 参数')
            sys.exit(1)
        register(args.name, args.master)
    else:
        parser.print_help()