app.py

import gradio as gr
import requests

# Fetch the flagged addresses once at startup from the data feed (static JSON file)
API_URL = "https://dt074px2e9qbh.cloudfront.net/exploit-api.json"
try:
    response = requests.get(API_URL)
    flagged_addresses = response.json()  # Load the JSON data into a list of strings or dictionaries
    # Ensure the data is a list of strings or dictionaries (addresses)
    if not isinstance(flagged_addresses, list):
        flagged_addresses = []
except Exception as e:
    flagged_addresses = []
    print(f"Failed to load API data: {e}")

# Function to infer the blockchain using Hugging Face API
def infer_chain(address, hf_token):
    headers = {
        "Authorization": f"Bearer {hf_token}"
    }
    prompt = f"Identify the blockchain of the following cryptocurrency address: {address}"

    # Make the request to Hugging Face's Inference API
    try:
        response = requests.post(
            "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1",
            headers=headers,
            json={"inputs": prompt}
        )
        response.raise_for_status()
        result = response.json()
        # Return the inferred blockchain (if available)
        return result[0]["generated_text"].strip()
    except Exception as e:
        return "Unknown"  # Fallback if there's an error

# Function to check if the address is flagged
def check_flagged_address(address: str, hf_token: str):
    """Check if a cryptocurrency address is flagged in the Bybit hack database."""
    
    # Normalize the input address (lowercase and stripped of whitespace)
    normalized_address = address.strip().lower()
    
    # Infer the chain of the address
    chain = infer_chain(address, hf_token)
    
    # Check if the address is in the flagged addresses
    flagged = False
    for entry in flagged_addresses:
        # Check if entry is a string (address)
        if isinstance(entry, str):
            if entry.strip().lower() == normalized_address:
                flagged = True
                break
        # Check if entry is a dictionary (with "address" key)
        elif isinstance(entry, dict) and "address" in entry:
            if entry["address"].strip().lower() == normalized_address:
                flagged = True
                break
    
    # Format the response
    if flagged:
        return f"Address {address} is Flagged ✅\nChain: {chain}\nStatus: Present in Bybit hack database"
    else:
        return f"Address {address} is Not Flagged ❌\nChain: {chain}\nStatus: Not present in Bybit hack database"

# Function to handle multiple addresses input
def handle_multiple_addresses(user_input, hf_token):
    """Process multiple addresses separated by commas."""
    results = []
    
    # Split input into multiple addresses by commas and strip any whitespace
    addresses = [addr.strip() for addr in user_input.split(",")]
    
    # Process each address individually
    for address in addresses:
        if address:  # Skip empty inputs
            result = check_flagged_address(address, hf_token)
            results.append(result)
        else:
            results.append("Empty address entered, please try again.")
    
    return "\n\n".join(results)  # Combine results with two newlines

# Gradio chat function
def chat_with_agent(user_input, chat_history, hf_token):
    if not chat_history:
        chat_history = []
    
    # If the user input is empty, respond like a chatbot
    if not user_input.strip():
        chat_history.append({"role": "user", "content": "I didn't enter an address."})
        chat_history.append({"role": "assistant", "content": "It seems like you didn't enter an address. Please input a valid cryptocurrency address to check if it’s flagged."})
        return chat_history, ""
    
    # Check if token is provided
    if not hf_token or hf_token.strip() == "":
        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": "Please enter a valid Hugging Face API token."})
        return chat_history, ""
    
    # Handle multiple addresses
    response = handle_multiple_addresses(user_input, hf_token)
    
    # Append as OpenAI-style messages
    chat_history.append({"role": "user", "content": user_input})
    chat_history.append({"role": "assistant", "content": response})
    
    # Return updated history and clear the input box
    return chat_history, ""

# Create the Gradio interface with tabs
with gr.Blocks(title="Bybit Hack Address Checker") as demo:
    gr.Markdown("# Bybit Hack Address Checker")
    
    with gr.Tabs():
        # Tab 1: Address Checker UI
        with gr.Tab(label="Check Address"):
            gr.Markdown("Enter one or more cryptocurrency addresses (separated by commas) to check if they're flagged in the Bybit hack database.")
            gr.Markdown("Provide your Hugging Face API token below (get it from [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)).")
            
            # HF Token input
            hf_token_input = gr.Textbox(
                placeholder="Enter your Hugging Face API token here",
                label="Hugging Face API Token",
                type="password",
                lines=1
            )
            
            # Chatbot component with messages format
            chatbot = gr.Chatbot(label="Conversation", type="messages")
            
            # Address input
            msg = gr.Textbox(
                placeholder="Enter address(es) here (separated by commas)",
                label="Address",
                lines=1
            )
            
            # Clear button
            clear = gr.Button("Clear")
            
            # Submit event
            msg.submit(
                fn=chat_with_agent,
                inputs=[msg, chatbot, hf_token_input],
                outputs=[chatbot, msg]
            )
            
            # Clear event
            clear.click(
                fn=lambda: ([], ""),
                inputs=None,
                outputs=[chatbot, msg]
            )
        
        # Tab 2: Explanation
        with gr.Tab(label="About This App"):
            gr.Markdown("""
                ## What This App Does
                This application helps you verify if a cryptocurrency address is associated with the Bybit hack, as tracked by the Elliptic dataset. Here’s how it works:

                1. **Input one or more Addresses**: Enter one or more wallet addresses (separated by commas).
                2. **Provide HF Token**: Supply a Hugging Face API token to access the AI model powering the agent.
                3. **Agent Processing**: The app uses a **language model** to infer the blockchain of each address (Ethereum, Bitcoin, Tron, etc.).
                4. **Flag Check**: It then checks each address against a list of flagged addresses from the Elliptic Bybit hack dataset (`https://dt074px2e9qbh.cloudfront.net/exploit-api.json`).
                5. **Result**: You’ll see whether each address is flagged ✅ or not ❌, its chain type, and if it’s present in the database.

                ### Why It’s Useful
                - **Security**: Quickly check if any address you're dealing with is linked to a known exploit.
                - **Free & Open**: Uses free-tier AI and public data, making it accessible to anyone with an HF account.
                - **Educational**: Demonstrates AI-driven blockchain analysis with a large language model.

                Get started by entering your HF token and addresses in the 'Check Address' tab!
            """)

# Launch the app
demo.launch()