Python script to monitor downloads folder and detect files with extensions maliciously tampered

fake_extension_detector.py

import os
import magic
from pydantic import BaseModel
from enum import Enum
import tkinter as tk
import time

# Tip: run it in the background with: nohup python fake_extension_detector.py &

SECONDS_TO_WAIT = 60
DIRECTORY = "/home/costa/Downloads"


class AnalysisStatus(Enum):
    SAFE = "SAFE"
    SUSPICIOUS = "SUSPICIOUS"
    NO_EXTENSION = "NO EXTENSION"
    UNKNOWN = "UNKNOWN"


class CustomBaseModel(BaseModel):
    class Config:
        frozen = True


class CheckFileReceived(CustomBaseModel):
    extension: str
    format: str


class CheckFileActual(CustomBaseModel):
    extension: str
    formats: list[str]

    def check(self, received: CheckFileReceived) -> bool:
        res = (received.format in self.formats) and (
            received.extension == self.extension
        )
        return res


def display_notification(
    text: str = "sample", analysis_status: AnalysisStatus = AnalysisStatus.SUSPICIOUS
):
    root = tk.Tk()

    red = "#b30000"
    orange = "#ff8000"
    background_color = "black"
    notification_title = "sample"
    text_color = "white"

    match analysis_status:
        case AnalysisStatus.SUSPICIOUS:
            background_color = red
            notification_title = "SUSPICIUS FILE FOUND"
        case AnalysisStatus.UNKNOWN:
            background_color = orange
            notification_title = "UNKNOWN FILE FOUND"

    root.title(notification_title)

    window_width = root.winfo_screenwidth() // 4
    window_height = root.winfo_screenheight() // 4

    x = (root.winfo_screenwidth() - window_width) // 2
    y = (root.winfo_screenheight() - window_height) // 2
    root.geometry(f"{window_width}x{window_height}+{x}+{y}")

    root.configure(bg=background_color)
    label = tk.Label(
        root, text=text, bg=background_color, fg=text_color, wraplength=400
    )

    label.place(relx=0.5, rely=0.5, anchor="center")
    root.mainloop()


def get_analysis_result(
    catalog: list[CheckFileActual], file: CheckFileReceived
) -> AnalysisStatus:
    for i in catalog:
        if file.extension not in [i.extension for i in catalog]:
            return AnalysisStatus.UNKNOWN
        if file.extension == "":
            return AnalysisStatus.NO_EXTENSION
        if i.check(file):
            return AnalysisStatus.SAFE
    return AnalysisStatus.SUSPICIOUS


def get_file_size_mb(file_path: str) -> float:
    return round(os.path.getsize(file_path) / (1024 * 1024), 2)


db3 = CheckFileActual(
    extension="db3",
    formats=[
        "SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 2, cookie 0x1, schema 4, UTF-8, version-valid-for 1"
    ],
)
pdf = CheckFileActual(extension="txt", formats=["ASCII text"])
py = CheckFileActual(extension="py", formats=["Python script, ASCII text executable"])
pdf = CheckFileActual(
    extension="pdf", formats=["PDF document, version 1.7 (zip deflate encoded)"]
)
wav = CheckFileActual(
    extension="wav", formats=["MPEG ADTS, layer III, v2,  32 kbps, 24 kHz, Monaural"]
)
deb = CheckFileActual(
    extension="deb",
    formats=[
        "Debian binary package (format 2.0), with control.tar.gz, data compression gz"
    ],
)
pem = CheckFileActual(extension="pem", formats=["OpenSSH private key"])
txt = CheckFileActual(extension="txt", formats=["ASCII text"])
toml = CheckFileActual(extension="toml", formats=["ASCII text"])
md = CheckFileActual(extension="md", formats=["ASCII text"])
sh = CheckFileActual(extension="sh", formats=["ASCII text"])
gitignore = CheckFileActual(extension="gitignore", formats=["ASCII text"])
out = CheckFileActual(extension="out", formats=["empty"])

CATALOG = [pdf, db3, py, pdf, wav, deb, pem, txt, toml, md, sh, gitignore, out]


def main():
    while True:
        file_list = [
            i
            for i in os.scandir(DIRECTORY)
            if i.is_file() and len(i.name.split(".")) > 1
        ]

        for i in file_list:
            text = None
            actual_type = magic.from_buffer(open(i.path, "rb").read(2048))
            claimed_type = i.name.split(".")[-1]

            received_file = CheckFileReceived(
                extension=claimed_type, format=actual_type
            )
            analysis_result = get_analysis_result(CATALOG, received_file)

            match analysis_result:
                case AnalysisStatus.SUSPICIOUS:
                    text = f"FILE: {i.name}\n\nACTUAL TYPE: {received_file.format}\n\nSIZE: {get_file_size_mb(i.path)} MB"
                case AnalysisStatus.UNKNOWN:
                    text = f"FILE: {i.name}\n\nTYPE: {received_file.format}\n\nSIZE: {get_file_size_mb(i.path)} MB\n\nINFO: You may want to update the extension catalog with this new filetype"
            if text:
                display_notification(text, analysis_result)

        time.sleep(SECONDS_TO_WAIT)


main()