Merge pull request #81 from R1c3P0T5/jerry/backend/all_on_zip

Jerry/backend/script and table

Author: Stanley1106

backend/README.md

@@ -22,15 +22,22 @@ uv run alembic upgrade head
 
 ## 匯入樣本（`scripts/load_malware_seeds.py`）
 
-1. 將含有惡意樣本的 zip 檔放到 `backend/data/samples/`
-   - Zip 預設密碼為 `infected`
-   - 單一 zip 可以包含多個樣本檔案
-2. （選用）在 `backend/data/malwareseed/` 內建立與 zip 同名的 JSON（`<zip_name>.json`）。
-   - 每個條目以 `sha256` 為主鍵，對應 zip 內樣本檔案的雜湊
-   - 未提供的欄位會自動套用預設值（platform = `Unknown`、family = `Unknown`、source = `ZIP:<zipname>`）
-   
+1. 將含有惡意樣本與對應 JSON 的 zip 放到 `backend/data/samples/`
+   - Zip 密碼為 `infected`
+2. JSON 條目以 `sha256` 作為主鍵，未填欄位會套用預設值
+3. JSON 檔名統一為 `seed.json`
 
-範例（`malwareseed/batch.json`）
+範例 zip 內容：
+
+```
+.zip
+└── seed.json    
+└── payload/
+    ├── sample-1.bin
+    └── sample-2.bin
+```
+
+範例 JSON：
 
 ```json
 [
@@ -44,27 +51,32 @@ uv run alembic upgrade head
 ]
 ```
 
-
-
-
 ## 匯入特徵並建立關聯（`scripts/load_feature_seeds.py`）
 
 1. 在 `backend/data/seed/` 放置特徵 JSON 檔
-2. 若需要將特徵綁定到特定樣本，於條目中新增 `sample_sha256` 欄位
-
-範例：
+2. 若需綁定樣本，於條目加入 `sample_sha256` 陣列
 
 ```json
-{
-  "feature_name": "",
-  "feature_type": "",
-  "yara_rule": "",
-  "sample_sha256": [
-    ""
-  ]
-}
+[
+  {
+    "feature_name": "",
+    "feature_type": "",
+    "yara_rule": null,
+    "machine_code": null,
+    "assembly": null,
+    "pseudo_c": "",
+    "behavior": "",
+    "tags": [
+      ""
+    ],
+    "created_by": "",
+    "mitre_ids": [
+      ""
+    ],
+    "sample_sha256": [
+      ""
+    ]
+  }
+]
 ```
 
-
-
-

backend/app/alembic/versions/32d263976ad4_change_family_datatype.py

@@ -0,0 +1,59 @@
+"""change family datatype
+
+Revision ID: 32d263976ad4
+Revises: af25782c1c4f
+Create Date: 2025-11-26 16:22:25.261198
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "32d263976ad4"
+down_revision: Union[str, None] = "af25782c1c4f"
+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.alter_column(
+        "malware",
+        "family",
+        existing_type=sa.VARCHAR(),
+        type_=sa.JSON(),
+        nullable=False,
+        server_default=sa.text("'[]'::json"),
+        postgresql_using="""
+        CASE
+            WHEN family IS NULL OR btrim(family) = '' THEN '[]'::json
+            WHEN family LIKE '[%%' THEN family::json
+            ELSE json_build_array(family)::json
+        END
+        """,
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column(
+        "malware",
+        "family",
+        existing_type=sa.JSON(),
+        type_=sa.VARCHAR(),
+        nullable=True,
+        server_default=None,
+        postgresql_using="""
+        CASE
+            WHEN family IS NULL OR json_array_length(family) = 0 THEN NULL
+            ELSE family->>0
+        END
+        """,
+    )
+    # ### end Alembic commands ###

backend/app/alembic/versions/af25782c1c4f_add_mitre_ids_for_malware_table.py

@@ -0,0 +1,37 @@
+"""add mitre_ids for malware table
+
+Revision ID: af25782c1c4f
+Revises: 962432c47268
+Create Date: 2025-11-26 14:31:49.003153
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "af25782c1c4f"
+down_revision: Union[str, None] = "962432c47268"
+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(
+        "malware",
+        sa.Column(
+            "mitre_ids", sa.JSON(), server_default=sa.text("'[]'::json"), nullable=False
+        ),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("malware", "mitre_ids")
+    # ### end Alembic commands ###

backend/app/malware/models.py

@@ -1,12 +1,14 @@
 from datetime import datetime
 from typing import TYPE_CHECKING
 
-from sqlalchemy import Column, Integer, text
+from pydantic import field_validator
+from sqlalchemy import JSON, Column, Integer, text
 from sqlalchemy.orm import relationship
 from sqlmodel import Field, Relationship, SQLModel
 
 from app.features.models import FeaturePublic
 from app.links.models import MalwareFeatureLink
+from app.malware.utils import normalize_family_values
 
 if TYPE_CHECKING:  # pragma: no cover - type checking only
     from app.features.models import Feature
@@ -19,9 +21,23 @@ class MalwareBase(SQLModel):
     md5: str
     size: int
     sha256: str
-    family: str | None = None
+    family: list[str] = Field(
+        default_factory=list,
+        sa_column=Column(JSON, nullable=False, server_default=text("'[]'::json")),
+        description="惡意家族標籤列表",
+    )
     source: str | None = None
     platform: str
+    mitre_ids: list[str] = Field(
+        default_factory=list,
+        sa_column=Column(JSON, nullable=False, server_default=text("'[]'::json")),
+        description="對應的 MITRE ATT&CK 技術 ID",
+    )
+
+    @field_validator("family", mode="before")
+    @classmethod
+    def _ensure_family_list(cls, value):  # type: ignore[override]
+        return normalize_family_values(value)
 
 
 class Malware(MalwareBase, table=True):
@@ -57,9 +73,17 @@ class MalwareUpdate(SQLModel, table=False):
 
     filename: str | None = None
     source: str | None = None
-    family: str | None = None
+    family: list[str] | None = None
     platform: str | None = None
     last_seen: datetime | None = None
+    mitre_ids: list[str] | None = None
+
+    @field_validator("family", mode="before")
+    @classmethod
+    def _ensure_update_family(cls, value):  # type: ignore[override]
+        if value is None:
+            return None
+        return normalize_family_values(value)
 
 
 class MalwareDelete(SQLModel):

backend/app/malware/routes.py

@@ -1,3 +1,4 @@
+import json
 import shutil
 from pathlib import Path
 from typing import List
@@ -31,7 +32,7 @@
     MalwarePublic,
     MalwareUpdate,
 )
-from .utils import calc_hashes_from_fileobj
+from .utils import calc_hashes_from_fileobj, normalize_family_values
 
 router = APIRouter(
     prefix="/malware", tags=["Malware"], dependencies=[Depends(get_current_user_dep)]
@@ -40,6 +41,26 @@
 SAMPLE_STORAGE_DIR = Path(__file__).resolve().parents[2] / "static" / "samples"
 
 
+def parse_mitre_ids(raw: str | None) -> list[str]:
+    if not raw:
+        return []
+    try:
+        parsed = json.loads(raw)
+    except json.JSONDecodeError:
+        parsed = None
+
+    values: list[str]
+    if isinstance(parsed, list):
+        values = [str(item).strip() for item in parsed if str(item).strip()]
+    else:
+        values = [segment.strip() for segment in raw.split(",") if segment.strip()]
+    return values
+
+
+def parse_family(raw: str | None) -> list[str]:
+    return normalize_family_values(raw)
+
+
 @router.get(
     "/{sample_id}",
     summary="根據 sample_id 取得惡意軟體資訊",
@@ -62,6 +83,7 @@ async def upload_malware(
     platform: str = Form(...),
     family: str | None = Form(None),
     source: str | None = Form(None),
+    mitre_ids: str | None = Form(None),
     db: Session = Depends(get_session),
 ):
     # 計算 hash
@@ -73,9 +95,10 @@ async def upload_malware(
         md5=md5,
         size=size,
         sha256=sha256,
-        family=family,
+        family=parse_family(family),
         source=source,
         platform=platform,
+        mitre_ids=parse_mitre_ids(mitre_ids),
     )
 
     if get_by_sha256(db, sha256):

backend/app/malware/utils.py

@@ -1,5 +1,43 @@
 import hashlib
-from typing import BinaryIO, Tuple
+from typing import Any, BinaryIO, Tuple
+
+
+def normalize_family_values(value: Any) -> list[str]:
+    """Normalize arbitrary input into a list of family labels."""
+
+    def _clean(items: list[str]) -> list[str]:
+        seen = set()
+        result: list[str] = []
+        for item in items:
+            cleaned = item.strip()
+            if not cleaned or cleaned in seen:
+                continue
+            seen.add(cleaned)
+            result.append(cleaned)
+        return result
+
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return _clean([str(item) for item in value])
+    if isinstance(value, str):
+        stripped = value.strip()
+        if not stripped:
+            return []
+        if stripped.startswith("[") and stripped.endswith("]"):
+            try:
+                import json
+
+                parsed = json.loads(stripped)
+            except json.JSONDecodeError:
+                parsed = None
+            if isinstance(parsed, list):
+                return _clean([str(item) for item in parsed])
+        fragments = [segment for segment in stripped.replace("\n", ",").split(",")]
+        return _clean(fragments) or _clean([stripped])
+    normalized = str(value).strip()
+    return _clean([normalized]) if normalized else []
+
 
 CHUNK = 1048576  # 1MB
 

backend/data/malwareseed/malware_samples.json

@@ -0,0 +1,9 @@
+[
+  {
+    "sha256": "c32fb305903a22106c6d3def0ac6c05b4f16cba99e23527b6c61d617ea794b1d",
+    "filename": "example.exe",
+    "platform": "windows",
+    "family": "Unknown",
+    "source": "MOTIF"
+  }
+]

backend/data/seed/malware_motif.json

@@ -16,6 +16,7 @@ dependencies = [
     "langchain>=1.0.0",
     "langchain-google-genai>=3.0.0",
     "pytest>=8.4.2",
+    "py7zr>=0.21.0",
 ]
 
 [dependency-groups]