mirror of
https://github.com/dobin/avred
synced 2026-06-08 13:54:13 +00:00
Dobin Rutishauser
140 lines
4.1 KiB
Python
140 lines
4.1 KiB
Python
import requests as req
|
|
import logging
|
|
import yara
|
|
import brotli
|
|
import os
|
|
import pickle
|
|
import hashlib
|
|
import time
|
|
import shutil
|
|
|
|
from model.model_base import Scanner
|
|
from config import config
|
|
|
|
|
|
HASHCACHE_FILE = "hashcache.pickle"
|
|
|
|
|
|
class HashCacheEntry():
|
|
def __init__(self, filename, result, scanTime, scannerName):
|
|
self.filename = filename
|
|
self.scanTime = scanTime
|
|
self.result = result
|
|
self.scannerName = scannerName
|
|
|
|
|
|
class HashCache():
|
|
def __init__(self):
|
|
self.cache = {}
|
|
|
|
def load(self):
|
|
if os.path.exists(HASHCACHE_FILE):
|
|
with open(HASHCACHE_FILE, "rb") as file:
|
|
logging.info("Loading HashCache")
|
|
self.cache = pickle.load(file)
|
|
logging.info(" {} hashes loaded".format(len(self.cache)) )
|
|
|
|
|
|
def save(self):
|
|
new = HASHCACHE_FILE + ".new"
|
|
with open(new, "wb") as file:
|
|
logging.info("Saving HashCache ({})".format(len(self.cache)))
|
|
pickle.dump(self.cache, file)
|
|
shutil.move(new, HASHCACHE_FILE)
|
|
|
|
|
|
def getResult(self, data, scannerName):
|
|
hash = hashlib.md5(data).hexdigest() + "_" + scannerName
|
|
|
|
if not hash in self.cache:
|
|
logging.debug("Not exist: {}".format(hash))
|
|
|
|
return self.cache.get(hash, None)
|
|
|
|
|
|
def addResult(self, data, filename, result, scanTime, scannerName):
|
|
hash = hashlib.md5(data).hexdigest() + "_" + scannerName
|
|
logging.debug("Add result: {}".format(hash))
|
|
self.cache[hash] = HashCacheEntry(filename, result, scanTime, scannerName)
|
|
|
|
|
|
hashCache = HashCache()
|
|
|
|
|
|
class ScannerRest(Scanner):
|
|
def __init__(self, url, name):
|
|
self.scanner_path = url
|
|
self.scanner_name = name
|
|
|
|
|
|
def scannerDetectsBytes(self, data: bytes, filename: str, useBrotli=True):
|
|
"""Returns true if file is detected"""
|
|
|
|
if config.get("hashCache") == True:
|
|
cacheResult = hashCache.getResult(data, self.scanner_name)
|
|
if cacheResult is not None:
|
|
return cacheResult.result
|
|
|
|
params = { 'filename': filename, 'brotli': useBrotli }
|
|
if useBrotli:
|
|
scanData = brotli.compress(data)
|
|
else:
|
|
scanData = data
|
|
|
|
timeStart = time.time()
|
|
try:
|
|
res = req.post(f"{self.scanner_path}/scan", params=params, data=scanData, timeout=10)
|
|
except:
|
|
# try again
|
|
logging.warning("Server timeout, retrying once")
|
|
logging.warning("Increase timeout if you scan large files")
|
|
res = req.post(f"{self.scanner_path}/scan", params=params, data=scanData, timeout=10)
|
|
jsonRes = res.json()
|
|
scanTime = round(time.time() - timeStart, 3)
|
|
|
|
# basically internal server error, e.g. AMSI not working
|
|
if res.status_code != 200:
|
|
logging.error("Error Code {}: {}".format(res.status_code, res.text))
|
|
raise Exception("Server error, aborting")
|
|
|
|
# server works and returns a result
|
|
ret_value = jsonRes['detected']
|
|
if config.get("hashCache") == True:
|
|
hashCache.addResult(data, filename, ret_value, scanTime, self.scanner_name)
|
|
return ret_value
|
|
|
|
|
|
def checkOnlineOrExit(self):
|
|
try:
|
|
res = req.post(f"{self.scanner_path}/test", timeout=1)
|
|
except:
|
|
logging.error("Scanner {} is not online at: {}".format(
|
|
self.scanner_name, self.scanner_path
|
|
))
|
|
exit(1)
|
|
|
|
|
|
class ScannerYara(Scanner):
|
|
def __init__(self, url, name):
|
|
self.scanner_path = url
|
|
self.scanner_name = name
|
|
|
|
|
|
def scannerDetectsBytes(self, data: bytes, filename: str):
|
|
"""Returns true if file is detected"""
|
|
rule = yara.compile(filepath=self.scanner_path)
|
|
matches = rule.match(data=data)
|
|
if len(matches) > 0:
|
|
return True
|
|
return False
|
|
|
|
|
|
def checkOnlineOrExit(self):
|
|
try:
|
|
rule = yara.compile(filepath=self.scanner_path)
|
|
except Exception as e:
|
|
logging.error("Scanner Yara failed for file {} error: {}".format(
|
|
self.scanner_path, e,
|
|
))
|
|
exit(1)
|