islab-malware-detector 0.5.0

Base framework for building malware detectors

islab-malware-detector

A base framework for building malware detectors with modern Python.

• Source code: https://github.com/bolin8017/islab-malware-detector.git
• Wiki: https://github.com/bolin8017/islab-malware-detector/wiki
• PyPI: https://pypi.org/project/islab-malware-detector/

Features

• Pydantic v2 Configuration - Type-safe config with env vars and file support
• Typer CLI - Extensible command-line interface via factory function
• Structured Logging - Console and JSON output formats with structlog
• Abstract Base Class - Define train, evaluate, and predict methods
• Type Hints - Full typing support throughout the codebase

Requirements

Tool	Version
Python	>= 3.12

Installation

pip install islab-malware-detector

Or with uv:

uv add islab-malware-detector

Quick Start

Basic Usage

from pathlib import Path
from typing import Any
from maldet import BaseDetector, BaseDetectorConfig

class MyDetector(BaseDetector):
    """My custom malware detector."""

    def train(self) -> Path:
        """Train the malware detection model."""
        self.logger.info("training_started", dataset=str(self.config.data.train))

        # Ensure output directory exists
        self.ensure_directory_exists(self.config.output.model)

        # Your training logic here
        # - Load training data from self.config.data.train
        # - Extract features and train model
        # - Save model to self.config.output.model

        self.logger.info("training_completed", model=str(self.config.output.model))
        return self.config.output.model

    def evaluate(self) -> dict[str, Any]:
        """Evaluate the trained model."""
        self.logger.info("evaluation_started", dataset=str(self.config.data.test))

        # Your evaluation logic here
        # - Load test data from self.config.data.test
        # - Load model from self.config.output.model
        # - Calculate metrics

        metrics = {
            "accuracy": 0.95,
            "precision": 0.93,
            "recall": 0.97,
            "f1_score": 0.95,
        }

        self.logger.info("evaluation_completed", **metrics)
        return metrics

    def predict(self) -> Path:
        """Run predictions on new data."""
        self.logger.info("prediction_started", input=str(self.config.data.predict))

        # Ensure output directory exists
        self.ensure_directory_exists(self.config.output.prediction)

        # Your prediction logic here
        # - Load input data from self.config.data.predict
        # - Load model from self.config.output.model
        # - Generate predictions
        # - Save results to self.config.output.prediction

        self.logger.info("prediction_completed", output=str(self.config.output.prediction))
        return self.config.output.prediction


# Create and use the detector
if __name__ == "__main__":
    detector = MyDetector()

    # Train the model
    model_path = detector.train()
    print(f"Model saved to: {model_path}")

    # Evaluate the model
    metrics = detector.evaluate()
    print(f"Evaluation metrics: {metrics}")

    # Run predictions
    predictions_path = detector.predict()
    print(f"Predictions saved to: {predictions_path}")

Configuration

Custom Configuration

You can extend the base configuration with your own fields:

from pydantic import Field
from pydantic_settings import SettingsConfigDict
from maldet import BaseDetectorConfig

class MyDetectorConfig(BaseDetectorConfig):
    """Custom configuration with additional fields."""

    model_config = SettingsConfigDict(
        env_prefix="MY_DETECTOR_",
        env_nested_delimiter="__",
    )

    # Custom fields for your detector
    batch_size: int = Field(default=32, description="Training batch size")
    model_name: str = Field(default="random_forest", description="Model type to use")
    use_gpu: bool = Field(default=False, description="Whether to use GPU acceleration")
    n_estimators: int = Field(default=100, description="Number of trees for random forest")


class MyDetector(BaseDetector):
    config_class = MyDetectorConfig

    def train(self) -> Path:
        self.logger.info(
            "training_config",
            batch_size=self.config.batch_size,
            model_name=self.config.model_name,
            use_gpu=self.config.use_gpu,
        )
        # Use self.config.batch_size, self.config.model_name, etc.
        ...

Environment Variables

Configure your detector using environment variables:

# Data paths
export MALWARE_DETECTOR_DATA__TRAIN="./data/train"
export MALWARE_DETECTOR_DATA__TEST="./data/test"
export MALWARE_DETECTOR_DATA__PREDICT="./data/predict"

# Output paths
export MALWARE_DETECTOR_OUTPUT__MODEL="./output/model.pkl"
export MALWARE_DETECTOR_OUTPUT__PREDICTION="./output/predictions.json"

# Logging
export MALWARE_DETECTOR_LOG__LEVEL="INFO"
export MALWARE_DETECTOR_LOG__FORMAT="console"

# Custom fields (for MyDetectorConfig)
export MY_DETECTOR_BATCH_SIZE=64
export MY_DETECTOR_MODEL_NAME="xgboost"
export MY_DETECTOR_USE_GPU=true

Configuration File

Save configuration as TOML file:

# config.toml
version = "0.3.0"

[data]
train = "./data/train"
test = "./data/test"
predict = "./data/predict"

[output]
model = "./output/model.pkl"
feature = "./output/features.pkl"
prediction = "./output/predictions.json"
log = "./logs/detector.log"

[log]
level = "INFO"
format = "console"

# Custom fields (for MyDetectorConfig)
batch_size = 64
model_name = "xgboost"
use_gpu = true
n_estimators = 200

Load configuration from file:

config = MyDetectorConfig.from_toml("config.toml")
detector = MyDetector(config=config)

CLI Integration

Create CLI for Your Detector

The framework provides automatic CLI generation:

# my_detector_cli.py
from maldet import BaseDetector, BaseDetectorConfig

class MyDetector(BaseDetector):
    # ... implementation ...
    pass

# Create CLI app
app = MyDetector.create_cli()

# Add custom commands
@app.command()
def info():
    """Display detector information."""
    print("My Malware Detector v1.0")
    print("Custom implementation using islab-malware-detector framework")

if __name__ == "__main__":
    app()

CLI Usage

# Display help
python my_detector_cli.py --help

# Train the model
python my_detector_cli.py train --config config.toml

# Evaluate the model
python my_detector_cli.py evaluate --config config.toml

# Run predictions
python my_detector_cli.py predict --config config.toml

# Generate default config file
python my_detector_cli.py init --output config.toml

# Custom logging
python my_detector_cli.py train --log-format json --log-level DEBUG

MLflow Tracking

Detectors auto-log to MLflow when MLFLOW_TRACKING_URI is set in the environment. When unset, behavior is unchanged — no MLflow dependency is required for normal use.

Install with the optional mlflow extra:

pip install "islab-malware-detector[mlflow]"

Environment variables:

Variable	Effect
MLFLOW_TRACKING_URI	When set, enables MLflow tracking
MLFLOW_RUN_ID	Reuse an existing run (platform creates it)

Logged artifacts per action:

• train: flattened config params, config.json, model directory under model/
• evaluate: numeric metrics, metrics.json (always written to output.log dir)
• predict: prediction file under prediction/

Example:

export MLFLOW_TRACKING_URI="http://mlflow.example.com"
export MLFLOW_RUN_ID="<run-id-created-by-platform>"

python my_detector_cli.py train --config config.json

Logging

The framework uses structlog for structured logging:

from maldet import configure_logging, get_logger

# Configure logging at application startup
configure_logging(level="INFO", format="console")

# Get a logger in your code
logger = get_logger(__name__)

# Log structured events
logger.info("model_training_started", dataset="train.csv", samples=1000)
logger.debug("feature_extracted", feature_count=128, elapsed_ms=45.2)
logger.warning("missing_data", file="sample.exe", reason="corrupted_header")
logger.error("training_failed", error_type="ValueError", message="Invalid input shape")

Output Formats

Console format (development):

2024-01-20T10:30:00 [info    ] model_training_started    dataset=train.csv samples=1000
2024-01-20T10:30:45 [debug   ] feature_extracted         feature_count=128 elapsed_ms=45.2

JSON format (production):

{"event": "model_training_started", "dataset": "train.csv", "samples": 1000, "timestamp": "2024-01-20T10:30:00", "level": "info"}
{"event": "feature_extracted", "feature_count": 128, "elapsed_ms": 45.2, "timestamp": "2024-01-20T10:30:45", "level": "debug"}

Architecture

Core Components

• BaseDetector: Abstract base class defining the interface (train, evaluate, predict)
• BaseDetectorConfig: Pydantic configuration with data paths, output paths, and logging
• CLI Factory: Automatic CLI generation with train/evaluate/predict commands
• Logging: Structured logging with console and JSON formats

Directory Structure

your-detector/
├── config.toml              # Configuration file
├── my_detector.py           # Your detector implementation
├── my_detector_cli.py       # CLI wrapper (optional)
├── data/
│   ├── train/               # Training data
│   ├── test/                # Test data
│   └── predict/             # Prediction input
├── output/
│   ├── model.pkl            # Trained model
│   ├── features.pkl         # Extracted features (optional)
│   └── predictions.json     # Prediction results
└── logs/
    └── detector.log         # Application logs

Migration from v0.2.x

Version 0.3.0 introduced breaking changes with a cleaner ABC-based design:

v0.2.x	v0.3.x
def stage_extract(self)	Removed - implement in train()
def stage_vectorize(self)	Removed - implement in train()
def stage_train(self)	def train(self) -> Path
def stage_predict(self)	def predict(self) -> Path
detector.run()	Call methods directly: train(), evaluate(), predict()
detector.run(stages=["extract"])	Not supported - call methods directly
default_stages class attribute	Not used - implement your own workflow

Migration example:

# v0.2.x (old)
class OldDetector(BaseDetector):
    def stage_extract(self):
        # Extract features
        return self.config.folder.feature

    def stage_train(self):
        # Train model
        return self.config.folder.model

    def stage_predict(self):
        # Predict
        return self.config.path.output

detector = OldDetector()
detector.run(stages=["extract", "train"])

# v0.3.x (new)
class NewDetector(BaseDetector):
    def train(self) -> Path:
        # Extract features + train model
        features = self._extract_features()
        model = self._train_model(features)
        return self.config.output.model

    def evaluate(self) -> dict[str, Any]:
        # Evaluate model
        return {"accuracy": 0.95}

    def predict(self) -> Path:
        # Predict
        predictions = self._run_inference()
        return self.config.output.prediction

detector = NewDetector()
detector.train()
detector.evaluate()
detector.predict()

Examples

See the examples directory for complete working examples:

• Basic Detector: Simple random forest classifier
• Deep Learning Detector: CNN-based detector with GPU support
• Ensemble Detector: Multiple models with voting

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

MIT License - see LICENSE.txt for details.

Citation

If you use this framework in your research, please cite:

@software{islab_malware_detector,
  title = {islab-malware-detector: A Framework for Building Malware Detectors},
  author = {PO-LIN LAI},
  year = {2024},
  url = {https://github.com/bolin8017/islab-malware-detector}
}