import os
import json
import requests
from dotenv import load_dotenv
import os
from time import perf_counter
from hashlib import sha256, sha1, md5
from colorama import Fore, Style
import vt
import json
# загружаем переменные окружения
load_dotenv()
# получаем API-ключ
VIRUS_TOTAL_API_KEY = os.getenv("VIRUS_TOTAL_KEY")
# Сигнатуры
signature_resources = ['res/signatures.txt', 'res/signatures2.txt',
'res/signatures3.txt', 'res/signatures4.txt']
# для своих требований, туда можно поместить хэши зловредов, которых нету в обычных списках сигнатур. А также описание и название.
signature_resource_info = 'res/signatures_info.json'
def add_signature(new_signature: dict) -> None:
"""Добавление сигнатуры в базу"""
with open(signature_resource_info, 'r') as file:
signatures_info = json.load(file)
for signature in new_signatures:
signatures_info['name'] = new_signatures['name']
signatures_info['desc'] = new_signatures['desc']
signatures_info['date'] = new_signatures['date']
with open(signature_resource_info, 'a') as file:
json.dump(signatures_info, file, indent=4)
return None
def rewrite_signatures(signatures: dict) -> bool:
"""Перезапись БД сигнатур"""
try:
with open(signature_resource_info, 'w') as file:
json.dump(signatures, file, indent=4)
except Exception as e:
return False
else:
return True
def get_info_signature(signature: str) -> str:
"""Получение информации об сигнатуре из БД"""
with open(signature_resource_info, 'r') as file:
signatures_info = json.load(file)
try:
signature_info = f'''Сигнатура: {signature}
Название угрозы: {signatures_info[signature]["name"]}
Описание угрозы: {signatures_info[signature]["desc"]}
Дата обнаружения угроза: {signatures_info[signature]["date"]}'''
except IndexError:
signature_info = f'По сигнатуре {signature} еще нету информации в наших базах данных'
finally:
return signature_info
def scan_file(filename: str, delete_file: bool, signature_resources: list, VT_API_KEY: str) -> None:
"""Сканируем файл
+ filename: str - путь до файла
+ delete_file: bool - удалять ли файл
+ signature_resources: list - список ресурсов сигнатур
+ VT_API_KEY: str - API ключ Virus Total"""
result = f'Сканируем {filename} на наличие угроз...\n'
try:
print('Load singatures...')
start = perf_counter()
shahash = sha256()
sha1hash = sha1()
md5hash = md5()
# Получаем хэши
with open(filename, 'rb') as file:
while True:
data = file.read()
if not data:
break
shahash.update(data) # sha256
sha1hash.update(data) # sha1
md5hash.update(data) #md5
result1 = shahash.hexdigest()
result2 = sha1hash.hexdigest()
result3 = md5hash.hexdigest()
result += f'Проверяем наличие хешей "{result1}", "{result2}", "{result3}" (sha256, sha1, md5) в сигнатурах...\n'
# Читаем сигнатуры
with open(signature_resources[0], 'r') as r:
signatures = list(r.read().split('\n'))
with open(signature_resources[1], 'r') as r:
for sign_hash2 in list(r.read().replace(';', '').split('\n')):
signatures.append(sign_hash2)
with open(signature_resources[2], 'r') as r:
for sign_hash3 in list(r.read().replace(';', '').split('\n')):
signatures.append(sign_hash3)
with open(signature_resources[3], 'r') as r:
for sign_hash4 in list(r.read().replace(';', '').split('\n')):
signatures.append(sign_hash4)
result += f'''Сканирование {filename} на Virus Total...\n'''
print('Loading Virus Total...')
client = vt.Client(VT_API_KEY)
# анализ файла
with open(filename, "rb") as f:
analysis = client.scan_file(f, wait_for_completion=True)
file = client.get_object(f"/files/{result1}")
stats = file.last_analysis_stats
result += f'VirusTotal: {filename}\t size: {file.size} bytes\n'
result += f'Безвредный: {stats["harmless"]}\n'
result += f'Неподдерживаемый тип: {stats["type-unsupported"]}\n'
result += f'Подозрительный: {stats["suspicious"]}\n'
result += f'Отказ: {stats["failure"]}\n'
result += f'Злонамеренный: {stats["malicious"]}\n'
result += f'Безопасный: {stats["undetected"]}\n'
print('Get final info...')
if result1 in signatures or result2 in signatures or result3 in signatures:
end = perf_counter()
total = end - start
result += f'Найдена угроза в файле {filename} (поиск по сигнатурам)\n'
# получаем информацию о файле, если он есть в сигнатурах
if result in signatures:
info = get_info_signature(result)
result += f'Информация: {info}\n'
elif result2 in signatures:
info = get_info_signature(result2)
result += f'Информация: {info}\n'
elif result3 in signatures:
info = get_info_signature(result3)
result += f'Информация: {info}\n'
# удаление файла
if delete_file:
os.remove(filename)
result += f'[!] Файл {filename} удален!\n'
result += f'Время работы: {(total):.07f}s\n'
else:
end = perf_counter()
total = end - start
result += 'Угроз не найдено\n'
result += f'Время работы: {(total):.07f}s\n'
except FileNotFoundError:
result += f'[!] Файл не найден\n'
except PermissionError:
result += f'[!] Ошибка прав доступа к файлу\n'
print('End')
return result
res = scan_file('<любой файл>', False, signature_resources, VIRUS_TOTAL_API_KEY)
print(res)