kk

app.py

@@ -1,7 +1,79 @@
+import torch
+from torch.nn.functional import softmax
+from transformers import GPT2Tokenizer
+import os
+import requests
+import tempfile
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Model loading and prediction function
+def check_vulnerabilities(solidity_file_path, model_directory='models/', device='cuda'):
+    device = 'cuda' if torch.cuda.is_available() and device == 'cuda' else 'cpu'
+    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+    tokenizer.pad_token = tokenizer.eos_token
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch(share=True)
+    with open(solidity_file_path, 'r', encoding='utf-8') as f:
+        test_code = f.read()
+
+    X = tokenizer([test_code], padding=True, truncation=True, return_tensors="pt")
+    dic_name = {
+        'reentrancy': 'Reentrancy Vulnerability',
+        'timestamp': 'Timestamp Dependency Vulnerability',
+        'delegatecall': 'Delegate Call Vulnerability',
+        'integeroverflow': 'Integer Overflow Vulnerability',
+    }
+    
+    dic01 = {0: 'The vulnerability does not exist', 1: 'The vulnerability exists'}
+    results = {}
+    
+    for model_name in os.listdir(model_directory):
+        vulnerability_name = dic_name[model_name.split('_')[0]]
+        cp_file = os.path.join(model_directory, model_name)
+        model = torch.load(cp_file, map_location=device)
+        X = X.to(device)
+        model.to(device)
+        model.eval()
+        pred = softmax(model(**X).logits, dim=1)[0]
+        results[vulnerability_name] = {
+            'result': dic01[int(pred.argmax(0))],
+            'confidence': pred.max().item()
+        }
+    
+    return results
+
+# Gradio interface function
+def check_vulnerabilities_interface(solidity_file_url, file_id):
+    model_directory = 'models/'
+    device = 'cuda'
+    
+    try:
+        response = requests.get(solidity_file_url)
+        response.raise_for_status()
+
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f"_{file_id}.sol") as temp_file:
+            temp_file.write(response.content)
+            temp_file.flush()
+            results = check_vulnerabilities(temp_file.name, model_directory, device)
+            os.remove(temp_file.name)
+        return results
+
+    except requests.exceptions.RequestException as e:
+        return {'error': f'Error fetching file: {e}'}
+    except Exception as e:
+        return {'error': str(e)}
+
+# Set up the Gradio interface
+interface = gr.Interface(
+    fn=check_vulnerabilities_interface,
+    inputs=[
+        gr.inputs.Textbox(label="Solidity File URL", placeholder="Enter URL here..."),
+        gr.inputs.Textbox(label="File ID", placeholder="Enter file ID here...")
+    ],
+    outputs="json",
+    title="Solidity Vulnerability Checker",
+    description="Enter the URL of a Solidity file and a file ID to check for vulnerabilities."
+)
+
+# Run the interface
+if __name__ == "__main__":
+    interface.launch()

requirements.txt

@@ -0,0 +1,68 @@
+appdirs==1.4.4
+asttokens==2.4.1
+certifi==2024.2.2
+charset-normalizer==3.3.2
+comm==0.2.2
+debugpy==1.8.1
+decorator==5.1.1
+exceptiongroup==1.2.0
+executing==2.0.1
+filelock==3.13.4
+fsspec==2024.3.1
+huggingface-hub==0.22.2
+idna==3.7
+ipykernel==6.29.4
+ipython==8.23.0
+jedi==0.19.1
+Jinja2==3.1.3
+jupyter_client==8.6.1
+jupyter_core==5.7.2
+Mako==1.3.3
+MarkupSafe==2.1.5
+matplotlib-inline==0.1.6
+mpmath==1.3.0
+nest-asyncio==1.6.0
+networkx==3.3
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.19.3
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.1.105
+packaging==24.0
+parso==0.8.4
+pexpect==4.9.0
+platformdirs==4.2.0
+prompt-toolkit==3.0.43
+psutil==5.9.8
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pycuda==2024.1
+Pygments==2.17.2
+python-dateutil==2.9.0.post0
+pytools==2024.1.1
+PyYAML==6.0.1
+pyzmq==25.1.2
+regex==2023.12.25
+requests==2.31.0
+safetensors==0.4.2
+six==1.16.0
+stack-data==0.6.3
+sympy==1.12
+tokenizers==0.15.2
+torch==2.2.2
+tornado==6.4
+tqdm==4.66.2
+traitlets==5.14.2
+transformers==4.39.3
+triton==2.2.0
+typing_extensions==4.11.0
+urllib3==2.2.1
+wcwidth==0.2.13