78 lines
No EOL
2.6 KiB
Python
78 lines
No EOL
2.6 KiB
Python
import os
|
|
import sys
|
|
import hashlib
|
|
import shutil
|
|
import math
|
|
from tqdm import tqdm
|
|
from sqlalchemy.orm import Session
|
|
import pandas as pd
|
|
from collections import Counter
|
|
from db.session import SessionLocal
|
|
from models.entry import Entry
|
|
from enums.file_type import FileType
|
|
from enums.type_enum import TypeEnum
|
|
|
|
df = pd.read_csv("./whitelist.csv", usecols=["sha256"])
|
|
sha_set = set(df["sha256"])
|
|
|
|
|
|
def calculate_entropy(file_path: str):
|
|
with open(file_path, "rb") as f:
|
|
data = f.read()
|
|
counts = Counter(data)
|
|
return -sum((c/len(data)) * math.log2(c/len(data)) for c in counts.values())
|
|
|
|
def get_hash(file_path: str, algorithm: str) -> str:
|
|
with open(file_path, "rb") as f:
|
|
return hashlib.file_digest(f, algorithm).hexdigest()
|
|
|
|
def main(folder_path: str, session: Session, file_len:int = -1, file_type_default: TypeEnum = TypeEnum.MALICIOUS):
|
|
files: list[str] = os.listdir(folder_path)
|
|
length = 0
|
|
if file_len == -1:
|
|
for file in files:
|
|
raise NotImplementedError()
|
|
else:
|
|
for i in tqdm(range(len(files))):
|
|
file_type = FileType.EXE
|
|
if len(files[i].split('.')) == 2:
|
|
if (files[i].split(".")[1] == "dll"):
|
|
file_type = FileType.DLL
|
|
file_path = f"{folder_path}/{files[i]}"
|
|
hash_sha256 = get_hash(file_path, "sha256")
|
|
if file_type_default is TypeEnum.MALICIOUS and hash_sha256 in sha_set:
|
|
continue
|
|
hash_md5 = get_hash(file_path, "md5")
|
|
hash_sha1 = get_hash(file_path, "sha1")
|
|
filename = hash_sha256
|
|
entropy = calculate_entropy(file_path)
|
|
shutil.copy(file_path, f"../export/bytes/{hash_sha256}")
|
|
entry = Entry(
|
|
filename = filename,
|
|
filename_actual = files[i],
|
|
md5 = hash_md5,
|
|
sha1 = hash_sha1,
|
|
sha256 = hash_sha256,
|
|
file_type = file_type,
|
|
entropy = entropy,
|
|
malicious = file_type_default
|
|
)
|
|
session.add(entry)
|
|
session.commit()
|
|
length += 1
|
|
if length == file_len:
|
|
break
|
|
|
|
|
|
if __name__ == "__main__":
|
|
session = SessionLocal()
|
|
if len(sys.argv) == 2:
|
|
main(sys.argv[1], session)
|
|
elif len(sys.argv) == 3:
|
|
main(sys.argv[1], session, int(sys.argv[2]))
|
|
elif len(sys.argv) == 4 and sys.argv[3] == "mal" or sys.argv[3] == "ben":
|
|
main(sys.argv[1], session, int(sys.argv[2]),
|
|
TypeEnum.BENIGN if sys.argv[3] == "ben" else TypeEnum.MALICIOUS
|
|
)
|
|
else:
|
|
raise ValueError() |