编辑代码

import requests
from concurrent.futures import ThreadPoolExecutor
import ipaddress
import netifaces
import argparse
import socket
from datetime import datetime

def get_local_network():
    """获取本机所在的局域网网段"""
    for interface in netifaces.interfaces():
        addrs = netifaces.ifaddresses(interface)
        if netifaces.AF_INET in addrs:
            for addr_info in addrs[netifaces.AF_INET]:
                ip = addr_info['addr']
                if ip == '127.0.0.1':
                    continue
                netmask = addr_info['netmask']
                # 计算 CIDR 表示法
                network = ipaddress.IPv4Network(f"{ip}/{netmask}", strict=False)
                return network
    return None

def scan_ollama(ip, port, timeout):
    """扫描指定 IP 和端口的 Ollama 服务"""
    url = f"http://{ip}:{port}/api/tags"
    try:
        response = requests.get(url, timeout=timeout)
        if response.status_code == 200:
            data = response.json()
            models = [model['name'] for model in data.get('models', [])]
            return (ip, port, models)
    except (requests.ConnectionError, requests.Timeout, requests.JSONDecodeError):
        pass
    return None

def main():
    parser = argparse.ArgumentParser(description='扫描局域网内的 Ollama 服务及模型')
    parser.add_argument('-p', '--ports', nargs='+', type=int, default=[11434],
                        help='要扫描的端口列表（默认：11434）')
    parser.add_argument('-t', '--threads', type=int, default=50,
                        help='并发线程数（默认：50）')
    parser.add_argument('-T', '--timeout', type=float, default=2.0,
                        help='请求超时时间（秒）（默认：2）')
    args = parser.parse_args()

    network = get_local_network()
    if not network:
        print("无法检测到本地网络接口")
        return

    print(f"扫描网络: {network}")
    print(f"扫描端口: {args.ports}")
    print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)

    found = []
    with ThreadPoolExecutor(max_workers=args.threads) as executor:
        futures = []
        for host in network.hosts():
            ip = str(host)
            for port in args.ports:
                futures.append(
                    executor.submit(scan_ollama, ip, port, args.timeout)
                )

        for future in futures:
            result = future.result()
            if result:
                found.append(result)

    if not found:
        print("未找到运行 Ollama 服务的主机")
        return

    print("\n发现以下 Ollama 实例：")
    for ip, port, models in found:
        print(f"IP: {ip}:{port}")
        print(f"可用模型 ({len(models)}):")
        for model in models:
            print(f"  - {model}")
        print("-" * 60)

if __name__ == "__main__":
    main()