Spaces:
Runtime error
Runtime error
File size: 2,937 Bytes
5e9aa8c 8e5d9cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import torch
from torch.nn.functional import softmax
from transformers import GPT2Tokenizer
import os
import requests
import tempfile
from flask import Flask, request, jsonify
import gradio as gr
app = Flask(__name__)
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
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)
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
@app.route('/check_vulnerabilities', methods=['POST'])
def check_vulnerabilities_api():
if 'solidity_file_url' not in request.json or 'file_id' not in request.json:
return jsonify({'error': 'No solidity_file_url or file_id provided'}), 400
solidity_file_url = request.json['solidity_file_url']
file_id = request.json['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 jsonify(results)
except requests.exceptions.RequestException as e:
return jsonify({'error': f'Error fetching file: {e}'}), 500
except Exception as e:
return jsonify({'error': str(e)}), 500
interface = gr.Interface(
fn=check_vulnerabilities,
inputs=[gr.inputs.Textbox(label="Solidity File URL"), gr.inputs.Textbox(label="File ID")],
outputs="json",
title="Solidity Vulnerability Checker",
description="Enter the URL of a Solidity file and a file ID to check for vulnerabilities."
)
if __name__ == "__main__":
interface.launch(share=True) |