import os
from ankml import ANKPredictor, ANKMLConfig

# 1. 基础配置
config = ANKMLConfig()
# 必须设置服务器地址
config.set('server_url', 'https://ankml.top') 

# 2. 初始化预测器
predictor = ANKPredictor(model_type='tall', config=config)

# 3. 定义恶意软件判定阈值(推荐0.9来消除一些误报。)
threshold = 0.9

# 4. 扫描文件
# 准备一个安全文件和恶意文件
safe_file = "path/to/your/safe_file.txt"
malicious_file = "path/to/your/malicious_file.exe"

files_to_scan = {
    "安全文件": safe_file,
    "恶意文件": malicious_file
}

for name, path in files_to_scan.items():
    if not os.path.exists(path):
        print(f"警告: {name} 未找到，跳过扫描。")
        continue

    print(f"\n正在扫描 {name}: {os.path.basename(path)}")
    try:
        result = predictor.predict(path)
        
        # 5. 解读结果
        probability = result.get('probability', 0.0)
        is_malware = probability >= threshold
        status = "恶意软件" if is_malware else "安全"

        print(f"  -> 状态: {status}")
        print(f"  -> 置信度: {probability:.4f}")

    except Exception as e:
        print(f"扫描失败: {e}")