84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
import json
|
|
import requests
|
|
import os
|
|
import hashlib
|
|
import modules.vars as horsy_vars
|
|
from rich import print
|
|
|
|
|
|
def add_to_cfg(key):
|
|
with open(horsy_vars.horsypath + 'config.cfg') as f:
|
|
config = json.load(f)
|
|
|
|
config['vt-key'] = key
|
|
|
|
with open(horsy_vars.horsypath + 'config.cfg', 'w') as f:
|
|
json.dump(config, f)
|
|
|
|
|
|
def get_key():
|
|
with open(horsy_vars.horsypath + 'config.cfg') as f:
|
|
config = json.load(f)
|
|
|
|
try:
|
|
return config['vt-key']
|
|
except KeyError:
|
|
return None
|
|
|
|
|
|
def scan_file(filename):
|
|
api_url = 'https://www.virustotal.com/api/v3/files'
|
|
headers = {'x-apikey': get_key()}
|
|
with open(filename, 'rb') as file:
|
|
files = {'file': (filename, file)}
|
|
if os.path.getsize(filename) < 33554432:
|
|
response = requests.post(api_url, headers=headers, files=files)
|
|
return response.json()['data']['id']
|
|
else:
|
|
api_url = 'https://www.virustotal.com/api/v3/files/upload_url'
|
|
response = requests.get(api_url, headers=headers)
|
|
response = requests.post(response.json()['data'], headers=headers, files=files)
|
|
return response.json()['data']['id']
|
|
|
|
|
|
def get_report(filename):
|
|
hash_md5 = hashlib.md5()
|
|
with open(filename, "rb") as f:
|
|
for chunk in iter(lambda: f.read(4096), b""):
|
|
hash_md5.update(chunk)
|
|
api_url = 'https://www.virustotal.com/api/v3/files/' + hash_md5.hexdigest()
|
|
headers = {'x-apikey': get_key()}
|
|
response = requests.get(api_url, headers=headers)
|
|
analysis = dict()
|
|
try:
|
|
analysis['detect'] = response.json()['data']['attributes']['last_analysis_stats']
|
|
except:
|
|
analysis['detect'] = 'No data'
|
|
try:
|
|
analysis['link'] = 'https://www.virustotal.com/gui/file/' + response.json()['data']['id']
|
|
except:
|
|
analysis['link'] = 'No data'
|
|
|
|
return analysis
|
|
|
|
|
|
def scan_to_cli(filename):
|
|
try:
|
|
print(f"Starting virustotal scan")
|
|
if not get_key():
|
|
print(f"[red]Virustotal api key not found[/]")
|
|
print(f"You can add it by entering [bold]horsy --vt \[your key][/] in terminal")
|
|
else:
|
|
print(f"[green]Virustotal api key found[/]")
|
|
print(f"[italic white]If you want to disable scan, type [/][bold]horsy --vt disable[/]"
|
|
f"[italic white] in terminal[/]")
|
|
scan_file(filename)
|
|
print(f"[green]Virustotal scan finished[/]")
|
|
analysis = get_report(filename)
|
|
print(f"[green]You can see report by opening: [white]{analysis['link']}[/]")
|
|
print(f"{analysis['detect']['malicious']} antivirus flagged this file as malicious")
|
|
|
|
return analysis
|
|
except:
|
|
return None
|