added images convert as well
This commit is contained in:
parent
cc99abee31
commit
5692f1bed0
13 changed files with 228 additions and 5 deletions
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -5,4 +5,6 @@
|
||||||
|
|
||||||
**/*.sh
|
**/*.sh
|
||||||
**/*.txt
|
**/*.txt
|
||||||
**/docker/**
|
**/docker/**
|
||||||
|
!**/docker/docker-compose.yml
|
||||||
|
!**/requirements.txt
|
||||||
1
.gitkeep
1
.gitkeep
|
|
@ -1 +0,0 @@
|
||||||
**/docker/docker-compose.yml
|
|
||||||
17
docker/docker-compose.yml
Normal file
17
docker/docker-compose.yml
Normal file
|
|
@ -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
|
||||||
30
images.py
Normal file
30
images.py
Normal file
|
|
@ -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()
|
||||||
4
main.py
4
main.py
|
|
@ -4,7 +4,7 @@ import hashlib
|
||||||
import shutil
|
import shutil
|
||||||
import math
|
import math
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
from sqlalchemy.orm import sessionmaker, Session
|
from sqlalchemy.orm import Session
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
from collections import Counter
|
from collections import Counter
|
||||||
from db.session import SessionLocal
|
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:
|
with open(file_path, "rb") as f:
|
||||||
return hashlib.file_digest(f, algorithm).hexdigest()
|
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)
|
files: list[str] = os.listdir(folder_path)
|
||||||
length = 0
|
length = 0
|
||||||
if file_len == -1:
|
if file_len == -1:
|
||||||
|
|
|
||||||
|
|
@ -26,6 +26,7 @@ if config.config_file_name is not None:
|
||||||
# target_metadata = mymodel.Base.metadata
|
# target_metadata = mymodel.Base.metadata
|
||||||
from db.base import Base
|
from db.base import Base
|
||||||
from models.entry import Entry # IMPORTANT: force model loading
|
from models.entry import Entry # IMPORTANT: force model loading
|
||||||
|
from models.images import Image
|
||||||
|
|
||||||
target_metadata = Base.metadata
|
target_metadata = Base.metadata
|
||||||
|
|
||||||
|
|
|
||||||
41
migrations/versions/2e92816ea461_.py
Normal file
41
migrations/versions/2e92816ea461_.py
Normal file
|
|
@ -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 ###
|
||||||
34
migrations/versions/376ae9f9cf1b_.py
Normal file
34
migrations/versions/376ae9f9cf1b_.py
Normal file
|
|
@ -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 ###
|
||||||
32
migrations/versions/3db63e16629b_.py
Normal file
32
migrations/versions/3db63e16629b_.py
Normal file
|
|
@ -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 ###
|
||||||
32
migrations/versions/f988852b4ab6_.py
Normal file
32
migrations/versions/f988852b4ab6_.py
Normal file
|
|
@ -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 ###
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
from enums.file_type import FileType
|
from enums.file_type import FileType
|
||||||
from enums.type_enum import TypeEnum
|
from enums.type_enum import TypeEnum
|
||||||
from sqlalchemy import Column, Integer, String, Enum, Double
|
from sqlalchemy import Column, Integer, String, Enum, Double
|
||||||
|
from sqlalchemy.orm import relationship
|
||||||
from db.base import Base
|
from db.base import Base
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -15,4 +16,4 @@ class Entry(Base):
|
||||||
sha256 = Column(String(255), nullable=False, unique=True)
|
sha256 = Column(String(255), nullable=False, unique=True)
|
||||||
file_type = Column(Enum(FileType), nullable=False)
|
file_type = Column(Enum(FileType), nullable=False)
|
||||||
entropy = Column(Double(), nullable=False)
|
entropy = Column(Double(), nullable=False)
|
||||||
malicious = Column(Enum(TypeEnum), nullable=False, name="type")
|
malicious = Column(Enum(TypeEnum), nullable=False, name="type")
|
||||||
|
|
|
||||||
13
models/images.py
Normal file
13
models/images.py
Normal file
|
|
@ -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)
|
||||||
21
requirements.txt
Normal file
21
requirements.txt
Normal file
|
|
@ -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
|
||||||
Loading…
Reference in a new issue