diff --git a/.gitignore b/.gitignore index 36460dd..112291e 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,6 @@ **/*.sh **/*.txt -**/docker/** \ No newline at end of file +**/docker/** +!**/docker/docker-compose.yml +!**/requirements.txt \ No newline at end of file diff --git a/.gitkeep b/.gitkeep deleted file mode 100644 index 4d871b3..0000000 --- a/.gitkeep +++ /dev/null @@ -1 +0,0 @@ -**/docker/docker-compose.yml \ No newline at end of file diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml new file mode 100644 index 0000000..e27f26d --- /dev/null +++ b/docker/docker-compose.yml @@ -0,0 +1,17 @@ +services: + db: + image: mariadb:latest + container_name: db + environment: + - MARIADB_ROOT_PASSWORD=1234 + ports: + - 3306:3306 + restart: unless-stopped + volumes: + - /home/laptop/dataset/scripts/docker/database:/var/lib/mysql + adminer: + image: adminer:latest + container_name: adminer + ports: + - 8080:8080 + restart: unless-stopped \ No newline at end of file diff --git a/images.py b/images.py new file mode 100644 index 0000000..25e2e72 --- /dev/null +++ b/images.py @@ -0,0 +1,30 @@ +from peconvert import convert, ImageTypes +from tqdm import tqdm +from sqlalchemy.orm import Session +import os +import sys +from db.session import SessionLocal +from models.entry import Entry +from models.images import Image + +def main(folder_path: str, export: str, session: Session): + files = os.listdir(folder_path) + for file in tqdm(files): + entry = session.query(Entry).filter(Entry.sha256 == file).first() + image = Image( + entry_id = entry.id, + filename = file + ) + file_path = f"{folder_path}/{file}" + export_path = f"{export}/{file}.png" + convert(file_path, export_path) + session.add(image) + session.commit() + + +if __name__ == "__main__": + session = SessionLocal() + if len(sys.argv) == 3: + main(sys.argv[1], sys.argv[2], session) + else: + raise ValueError() \ No newline at end of file diff --git a/main.py b/main.py index cdc7b64..ae14ffb 100644 --- a/main.py +++ b/main.py @@ -4,7 +4,7 @@ import hashlib import shutil import math from tqdm import tqdm -from sqlalchemy.orm import sessionmaker, Session +from sqlalchemy.orm import Session import pandas as pd from collections import Counter from db.session import SessionLocal @@ -26,7 +26,7 @@ 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): +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: diff --git a/migrations/env.py b/migrations/env.py index 2418eae..f8c13d2 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -26,6 +26,7 @@ if config.config_file_name is not None: # target_metadata = mymodel.Base.metadata from db.base import Base from models.entry import Entry # IMPORTANT: force model loading +from models.images import Image target_metadata = Base.metadata diff --git a/migrations/versions/2e92816ea461_.py b/migrations/versions/2e92816ea461_.py new file mode 100644 index 0000000..36e2bdc --- /dev/null +++ b/migrations/versions/2e92816ea461_.py @@ -0,0 +1,41 @@ +""" + +Revision ID: 2e92816ea461 +Revises: 36eb607a5bf0 +Create Date: 2026-06-14 19:54:48.977452 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '2e92816ea461' +down_revision: Union[str, Sequence[str], None] = '36eb607a5bf0' +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('images', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('entry_id', sa.Integer(), nullable=False), + sa.Column('path', sa.String(length=255), nullable=False), + sa.ForeignKeyConstraint(['entry_id'], ['entries.id'], ), + sa.PrimaryKeyConstraint('id'), + sa.UniqueConstraint('entry_id') + ) + op.create_index(op.f('ix_images_id'), 'images', ['id'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f('ix_images_id'), table_name='images') + op.drop_table('images') + # ### end Alembic commands ### diff --git a/migrations/versions/376ae9f9cf1b_.py b/migrations/versions/376ae9f9cf1b_.py new file mode 100644 index 0000000..eb866d5 --- /dev/null +++ b/migrations/versions/376ae9f9cf1b_.py @@ -0,0 +1,34 @@ +""" + +Revision ID: 376ae9f9cf1b +Revises: 2e92816ea461 +Create Date: 2026-06-14 19:55:33.468092 + +""" +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 = '376ae9f9cf1b' +down_revision: Union[str, Sequence[str], None] = '2e92816ea461' +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('images', sa.Column('filename', sa.String(length=255), nullable=False)) + op.drop_column('images', 'path') + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('images', sa.Column('path', mysql.VARCHAR(length=255), nullable=False)) + op.drop_column('images', 'filename') + # ### end Alembic commands ### diff --git a/migrations/versions/3db63e16629b_.py b/migrations/versions/3db63e16629b_.py new file mode 100644 index 0000000..79b5a31 --- /dev/null +++ b/migrations/versions/3db63e16629b_.py @@ -0,0 +1,32 @@ +""" + +Revision ID: 3db63e16629b +Revises: 376ae9f9cf1b +Create Date: 2026-06-14 20:16:44.750486 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '3db63e16629b' +down_revision: Union[str, Sequence[str], None] = '376ae9f9cf1b' +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 ### diff --git a/migrations/versions/f988852b4ab6_.py b/migrations/versions/f988852b4ab6_.py new file mode 100644 index 0000000..6752b22 --- /dev/null +++ b/migrations/versions/f988852b4ab6_.py @@ -0,0 +1,32 @@ +""" + +Revision ID: f988852b4ab6 +Revises: 3db63e16629b +Create Date: 2026-06-14 20:17:26.165235 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'f988852b4ab6' +down_revision: Union[str, Sequence[str], None] = '3db63e16629b' +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 ### diff --git a/models/entry.py b/models/entry.py index d0d693d..9ba60ab 100644 --- a/models/entry.py +++ b/models/entry.py @@ -1,6 +1,7 @@ from enums.file_type import FileType from enums.type_enum import TypeEnum from sqlalchemy import Column, Integer, String, Enum, Double +from sqlalchemy.orm import relationship from db.base import Base @@ -15,4 +16,4 @@ class Entry(Base): 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") \ No newline at end of file + malicious = Column(Enum(TypeEnum), nullable=False, name="type") diff --git a/models/images.py b/models/images.py new file mode 100644 index 0000000..d87e644 --- /dev/null +++ b/models/images.py @@ -0,0 +1,13 @@ +from sqlalchemy import Column, Integer, String, ForeignKey +from sqlalchemy.orm import relationship +from db.base import Base + + +class Image(Base): + __tablename__ = "images" + + id = Column(Integer, primary_key=True, index=True) + entry_id = Column(Integer, ForeignKey("entries.id"), unique=True, nullable=False) + + # example fields for image + filename = Column(String(255), nullable=False) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..54da8da --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +alembic==1.18.4 +contourpy==1.3.3 +cycler==0.12.1 +fonttools==4.63.0 +greenlet==3.5.1 +kiwisolver==1.5.0 +Mako==1.3.12 +MarkupSafe==3.0.3 +matplotlib==3.11.0 +numpy==2.4.6 +packaging==26.2 +pandas==3.0.3 +peconvert @ git+https://github.com/vmillios/peconvert.git@8687d58117ead8e2fc473952ed6e15b81a1ab8ee +pillow==12.2.0 +PyMySQL==1.2.0 +pyparsing==3.3.2 +python-dateutil==2.9.0.post0 +six==1.17.0 +SQLAlchemy==2.0.50 +tqdm==4.68.2 +typing_extensions==4.15.0