211 lines
8.7 KiB
Python
211 lines
8.7 KiB
Python
import os, yara, numpy, onnxruntime
|
|
import ctypes, ctypes.wintypes, pefile
|
|
from PIL import Image
|
|
|
|
####################################################################################################
|
|
|
|
class GUID(ctypes.Structure):
|
|
_fields_ = [
|
|
("Data1", ctypes.wintypes.DWORD),
|
|
("Data2", ctypes.wintypes.WORD),
|
|
("Data3", ctypes.wintypes.WORD),
|
|
("Data4", ctypes.c_ubyte * 8)]
|
|
|
|
class WINTRUST_FILE_INFO(ctypes.Structure):
|
|
_fields_ = [
|
|
("cbStruct", ctypes.wintypes.DWORD),
|
|
("pcwszFilePath", ctypes.wintypes.LPCWSTR),
|
|
("hFile", ctypes.wintypes.HANDLE),
|
|
("pgKnownSubject", ctypes.wintypes.LPVOID)]
|
|
|
|
class WINTRUST_DATA(ctypes.Structure):
|
|
_fields_ = [
|
|
("cbStruct", ctypes.wintypes.DWORD),
|
|
("pPolicyCallbackData", ctypes.wintypes.LPVOID),
|
|
("pSIPClientData", ctypes.wintypes.LPVOID),
|
|
("dwUIChoice", ctypes.wintypes.DWORD),
|
|
("fdwRevocationChecks", ctypes.wintypes.DWORD),
|
|
("dwUnionChoice", ctypes.wintypes.DWORD),
|
|
("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)),
|
|
("dwStateAction", ctypes.wintypes.DWORD),
|
|
("hWVTStateData", ctypes.wintypes.HANDLE),
|
|
("pwszURLReference", ctypes.wintypes.LPCWSTR),
|
|
("dwProvFlags", ctypes.wintypes.DWORD),
|
|
("dwUIContext", ctypes.wintypes.DWORD),
|
|
("pSignatureSettings", ctypes.wintypes.LPVOID)]
|
|
|
|
####################################################################################################
|
|
|
|
class sign_scanner:
|
|
def __init__(self):
|
|
self.verify = GUID(0x00AAC56B, 0xCD44, 0x11D0, (0x8C, 0xC2, 0x00, 0xC0, 0x4F, 0xC2, 0x95, 0xEE))
|
|
|
|
def init_windll(self, path):
|
|
for name in path:
|
|
try:
|
|
setattr(self, name.lower(), ctypes.WinDLL(name, use_last_error=True))
|
|
except Exception as e:
|
|
self.send_message(e, "warn", False)
|
|
|
|
self.WinVerifyTrust = self.wintrust.WinVerifyTrust
|
|
self.WinVerifyTrust.restype = ctypes.wintypes.LONG
|
|
self.WinVerifyTrust.argtypes = [ctypes.wintypes.HWND, ctypes.POINTER(GUID), ctypes.c_void_p]
|
|
|
|
def is_sign(self, file_path):
|
|
with pefile.PE(file, fast_load=True) as pe:
|
|
sec_dir = pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_SECURITY']]
|
|
return sec_dir.VirtualAddress != 0 and sec_dir.Size > 0
|
|
|
|
def sign_verify(self, file_path):
|
|
fi = WINTRUST_FILE_INFO(ctypes.sizeof(WINTRUST_FILE_INFO), file_path, None, None)
|
|
data = WINTRUST_DATA(ctypes.sizeof(WINTRUST_DATA), None, None, 2, 0, 1,
|
|
ctypes.pointer(fi), 1, None, None, 0, 0, None)
|
|
s = self.WinVerifyTrust(None, ctypes.byref(self.verify), ctypes.byref(data))
|
|
data.dwStateAction = 2
|
|
self.WinVerifyTrust(None, ctypes.byref(self.verify), ctypes.byref(data))
|
|
return s == 0
|
|
|
|
####################################################################################################
|
|
|
|
class rule_scanner:
|
|
def __init__(self):
|
|
self.rules = {}
|
|
self.network = []
|
|
|
|
def load_path(self, path):
|
|
for root, _, files in os.walk(path):
|
|
for file in files:
|
|
self.load_file(os.path.join(root, file))
|
|
|
|
def load_file(self, file):
|
|
ext = os.path.splitext(file)[1].lower()
|
|
try:
|
|
if ext in ('.yara', '.yar'):
|
|
self.rules[file] = yara.compile(file)
|
|
elif ext in ('.yc', '.yrc'):
|
|
self.rules[file] = yara.load(file)
|
|
elif ext in ('.ip', '.txt'):
|
|
with open(file, "r", encoding="utf-8", errors="ignore") as f:
|
|
self.network.extend(line.strip() for line in f if line.strip())
|
|
except Exception:
|
|
return False
|
|
|
|
def yara_scan(self, file_path):
|
|
try:
|
|
for rules in self.rules.values():
|
|
matches = rules.match(filepath=file_path)
|
|
if matches:
|
|
rule_name = str(matches[0])
|
|
label = rule_name.split("_")[0]
|
|
level = rule_name.split("_")[-1]
|
|
return f"Rules/{label}", level
|
|
return False, False
|
|
except Exception:
|
|
return False, False
|
|
|
|
####################################################################################################
|
|
|
|
class model_scanner:
|
|
def __init__(self):
|
|
self.models = {}
|
|
self.suffix = {
|
|
".com", ".dll", ".drv", ".exe", ".ocx", ".scr", ".sys",
|
|
".bat", ".cmd", ".js", ".php", ".ps1", ".vbs", ".wsf"}
|
|
self.labels = [
|
|
"Pefile/White", "Script/White", "Pefile/General", "Script/General"]
|
|
self.detect = {
|
|
"Pefile/General", "Script/General"}
|
|
self.resize = (224, 224)
|
|
|
|
def load_path(self, path):
|
|
for root, _, files in os.walk(path):
|
|
for file in files:
|
|
self.load_file(os.path.join(root, file))
|
|
|
|
def load_file(self, file):
|
|
if file.lower().endswith('.onnx'):
|
|
try:
|
|
self.models[file] = onnxruntime.InferenceSession(file)
|
|
except Exception:
|
|
return False
|
|
|
|
def model_scan(self, file_path):
|
|
try:
|
|
model = next(iter(self.models.values()))
|
|
shape = model.get_inputs()[0].shape
|
|
channels = shape[-1] if shape[-1] in (1, 3) else (shape[1] if shape[1] in (1, 3) else 1)
|
|
|
|
file_sections = self.get_type(file_path)
|
|
if not file_sections:
|
|
return False, False
|
|
|
|
batch_images = [self.preprocess_image(data, self.resize, channels) for data in file_sections.values()]
|
|
arr = numpy.stack([numpy.asarray(img).astype('float32') / 255.0 for img in batch_images])
|
|
if arr.ndim == 3:
|
|
arr = arr[..., None]
|
|
if shape[-1] in (1, 3):
|
|
model_input = arr
|
|
elif shape[1] in (1, 3):
|
|
model_input = arr.transpose(0, 3, 1, 2)
|
|
else:
|
|
return False, False
|
|
|
|
input_name = model.get_inputs()[0].name
|
|
predictions = model.run(None, {input_name: model_input})[0]
|
|
|
|
label_scores = {label: [] for label in self.labels}
|
|
for pre_answers in predictions:
|
|
for k, score in enumerate(pre_answers):
|
|
label_scores[self.labels[k]].append(score)
|
|
label_percentage = {label: (sum(scores)/len(scores))*100 for label, scores in label_scores.items() if scores}
|
|
if not label_percentage:
|
|
return False, False
|
|
best_label, best_conf = max(label_percentage.items(), key=lambda x: x[1])
|
|
if best_label in self.detect:
|
|
return best_label, int(best_conf)
|
|
return False, False
|
|
except Exception:
|
|
return False, False
|
|
|
|
def preprocess_image(self, file_data, size, channels=1):
|
|
width, height = size
|
|
wah = int(numpy.ceil(numpy.sqrt(len(file_data) / channels)))
|
|
arr = numpy.frombuffer(file_data, dtype=numpy.uint8)
|
|
img = numpy.zeros(wah*wah*channels, dtype=numpy.uint8)
|
|
img[:len(file_data)] = arr
|
|
|
|
if channels == 1:
|
|
image = Image.fromarray(img.reshape((wah, wah)), 'L')
|
|
else:
|
|
image = Image.fromarray(img.reshape((wah, wah, channels)), 'RGB')
|
|
return image.resize((width, height), Image.Resampling.NEAREST)
|
|
|
|
def is_text_file(self, content, sample_size=1024):
|
|
raw = content[:sample_size]
|
|
if not raw:
|
|
return False
|
|
text_char = set(range(32, 127)) | {9, 10, 13}
|
|
nontext = sum(b not in text_char for b in raw)
|
|
return nontext / len(raw) < 0.15
|
|
|
|
def get_type(self, file_path):
|
|
match_data = {}
|
|
if not os.path.splitext(file_path)[-1].lower() in self.suffix:
|
|
return match_data
|
|
try:
|
|
with pefile.PE(file_path, fast_load=True) as pe:
|
|
sec = pe.OPTIONAL_HEADER.DATA_DIRECTORY[pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_SECURITY']]
|
|
if bool(sec.VirtualAddress and sec.Size):
|
|
return {}
|
|
for section in pe.sections:
|
|
name = section.Name.rstrip(b'\x00').decode('latin1').lower()
|
|
if section.Characteristics & 0x00000020:
|
|
match_data[name] = section.get_data()
|
|
|
|
except pefile.PEFormatError:
|
|
with open(file_path, 'rb') as f:
|
|
file_content = f.read()
|
|
if self.is_text_file(file_content):
|
|
match_data[os.path.splitext(file_path)[-1].lower()] = file_content
|
|
return match_data
|