thesis-scripts/main.py

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()