Initial Commit

This commit is contained in:
Vasileios Millios 2026-06-14 19:38:09 +00:00
commit cc99abee31
24 changed files with 87338 additions and 0 deletions

8
.gitignore vendored Normal file
View file

@ -0,0 +1,8 @@
**/.venv
**/__pycache__/**
**/__pycache__
**/*.pyc
**/*.sh
**/*.txt
**/docker/**

1
.gitkeep Normal file
View file

@ -0,0 +1 @@
**/docker/docker-compose.yml

35
alembic.ini Normal file
View file

@ -0,0 +1,35 @@
[alembic]
script_location = migrations
sqlalchemy.url = mysql+pymysql://root:1234@localhost:3306/thesis
[loggers]
keys = root,sqlalchemy,alembic
[handlers]
keys = console
[formatters]
keys = generic
[logger_root]
level = WARN
handlers = console
[logger_sqlalchemy]
level = WARN
handlers =
qualname = sqlalchemy.engine
[logger_alembic]
level = INFO
handlers =
qualname = alembic
[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic
[formatter_generic]
format = %(levelname)-5.5s [%(name)s] %(message)s

0
db/__init__.py Normal file
View file

4
db/base.py Normal file
View file

@ -0,0 +1,4 @@
from sqlalchemy.orm import DeclarativeBase
class Base(DeclarativeBase):
pass

5
db/engine.py Normal file
View file

@ -0,0 +1,5 @@
from sqlalchemy import create_engine
engine = create_engine(
"mysql+pymysql://root:1234@localhost:3306/thesis"
)

8
db/session.py Normal file
View file

@ -0,0 +1,8 @@
from sqlalchemy.orm import sessionmaker
from db.engine import engine
SessionLocal = sessionmaker(
bind=engine,
autoflush=False,
autocommit=False
)

0
enums/__init.py Normal file
View file

5
enums/file_type.py Normal file
View file

@ -0,0 +1,5 @@
from enum import Enum, auto
class FileType(Enum):
EXE = auto()
DLL = auto()

5
enums/type_enum.py Normal file
View file

@ -0,0 +1,5 @@
from enum import Enum, auto
class TypeEnum(Enum):
MALICIOUS = auto()
BENIGN = auto()

78
main.py Normal file
View file

@ -0,0 +1,78 @@
import os
import sys
import hashlib
import shutil
import math
from tqdm import tqdm
from sqlalchemy.orm import sessionmaker, 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: sessionmaker[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()

1
migrations/README Normal file
View file

@ -0,0 +1 @@
Generic single-database configuration.

87
migrations/env.py Normal file
View file

@ -0,0 +1,87 @@
import sys
from pathlib import Path
BASE_DIR = Path(__file__).resolve().parents[1]
sys.path.append(str(BASE_DIR))
from logging.config import fileConfig
from sqlalchemy import engine_from_config
from sqlalchemy import pool
from alembic import context
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config
# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
fileConfig(config.config_file_name)
# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
from db.base import Base
from models.entry import Entry # IMPORTANT: force model loading
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.
def run_migrations_offline() -> None:
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
url = config.get_main_option("sqlalchemy.url")
context.configure(
url=url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)
with context.begin_transaction():
context.run_migrations()
def run_migrations_online() -> None:
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
connectable = engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
with connectable.connect() as connection:
context.configure(
connection=connection, target_metadata=target_metadata
)
with context.begin_transaction():
context.run_migrations()
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()

28
migrations/script.py.mako Normal file
View file

@ -0,0 +1,28 @@
"""${message}
Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
${imports if imports else ""}
# revision identifiers, used by Alembic.
revision: str = ${repr(up_revision)}
down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
def upgrade() -> None:
"""Upgrade schema."""
${upgrades if upgrades else "pass"}
def downgrade() -> None:
"""Downgrade schema."""
${downgrades if downgrades else "pass"}

View file

@ -0,0 +1,32 @@
"""initial add
Revision ID: 095797896701
Revises: 9071255c7ee9
Create Date: 2026-06-14 16:30:54.744028
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '095797896701'
down_revision: Union[str, Sequence[str], None] = '9071255c7ee9'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
pass
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
pass
# ### end Alembic commands ###

View file

@ -0,0 +1,38 @@
"""add content to entry
Revision ID: 16d350f2e166
Revises:
Create Date: 2026-06-14 16:20:09.110465
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '16d350f2e166'
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('entries',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('title', sa.String(length=255), nullable=False),
sa.PrimaryKeyConstraint('id')
)
op.create_index(op.f('ix_entries_id'), 'entries', ['id'], unique=False)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index(op.f('ix_entries_id'), table_name='entries')
op.drop_table('entries')
# ### end Alembic commands ###

View file

@ -0,0 +1,44 @@
"""initial add
Revision ID: 271c3e6de3d7
Revises: 16d350f2e166
Create Date: 2026-06-14 16:27:43.274217
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import mysql
# revision identifiers, used by Alembic.
revision: str = '271c3e6de3d7'
down_revision: Union[str, Sequence[str], None] = '16d350f2e166'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('entries', sa.Column('filename', sa.String(length=255), nullable=False))
op.add_column('entries', sa.Column('sha1', sa.String(length=255), nullable=False))
op.add_column('entries', sa.Column('sha256', sa.String(length=255), nullable=False))
op.add_column('entries', sa.Column('file_type', sa.Enum('EXE', 'DLL', name='filetype'), nullable=False))
op.add_column('entries', sa.Column('entropy', sa.Double(), nullable=False))
op.add_column('entries', sa.Column('type', sa.Enum('MALICIOUS', 'BENIGN', name='type'), nullable=False))
op.drop_column('entries', 'title')
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('entries', sa.Column('title', mysql.VARCHAR(length=255), nullable=False))
op.drop_column('entries', 'type')
op.drop_column('entries', 'entropy')
op.drop_column('entries', 'file_type')
op.drop_column('entries', 'sha256')
op.drop_column('entries', 'sha1')
op.drop_column('entries', 'filename')
# ### end Alembic commands ###

View file

@ -0,0 +1,32 @@
"""
Revision ID: 36eb607a5bf0
Revises: 89b6c7607e31
Create Date: 2026-06-14 17:27:22.739565
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '36eb607a5bf0'
down_revision: Union[str, Sequence[str], None] = '89b6c7607e31'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('entries', sa.Column('filename_actual', sa.String(length=255), nullable=True))
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column('entries', 'filename_actual')
# ### end Alembic commands ###

View file

@ -0,0 +1,32 @@
"""sha256 unique
Revision ID: 78ead4bb04a8
Revises: 095797896701
Create Date: 2026-06-14 16:44:01.073200
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '78ead4bb04a8'
down_revision: Union[str, Sequence[str], None] = '095797896701'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.create_unique_constraint(None, 'entries', ['sha256'])
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(None, 'entries', type_='unique')
# ### end Alembic commands ###

View file

@ -0,0 +1,32 @@
"""update
Revision ID: 89b6c7607e31
Revises: 78ead4bb04a8
Create Date: 2026-06-14 16:46:37.313636
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '89b6c7607e31'
down_revision: Union[str, Sequence[str], None] = '78ead4bb04a8'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
pass
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
pass
# ### end Alembic commands ###

View file

@ -0,0 +1,32 @@
"""initial add
Revision ID: 9071255c7ee9
Revises: 271c3e6de3d7
Create Date: 2026-06-14 16:28:13.879143
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '9071255c7ee9'
down_revision: Union[str, Sequence[str], None] = '271c3e6de3d7'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('entries', sa.Column('md5', sa.String(length=255), nullable=False))
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column('entries', 'md5')
# ### end Alembic commands ###

0
models/__init__.py Normal file
View file

18
models/entry.py Normal file
View file

@ -0,0 +1,18 @@
from enums.file_type import FileType
from enums.type_enum import TypeEnum
from sqlalchemy import Column, Integer, String, Enum, Double
from db.base import Base
class Entry(Base):
__tablename__ = "entries"
id = Column(Integer, primary_key=True, index=True, )
filename = Column(String(255), nullable=False)
filename_actual = Column(String(255))
md5 = Column(String(255), nullable=False)
sha1 = Column(String(255), nullable=False)
sha256 = Column(String(255), nullable=False, unique=True)
file_type = Column(Enum(FileType), nullable=False)
entropy = Column(Double(), nullable=False)
malicious = Column(Enum(TypeEnum), nullable=False, name="type")

86813
whitelist.csv Normal file

File diff suppressed because it is too large Load diff