islab-malware-detector 0.4.0
Base framework for building malware detectors
License: MIT
Requires Python: >=3.12
Keywords: detector, framework, malware, security

# 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

```bash
pip install islab-malware-detector
```

Or with [uv](https://github.com/astral-sh/uv):

```bash
uv add islab-malware-detector
```

## Quick Start

### Basic Usage

```python
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:

```python
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:

```bash
# 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:

```toml
# 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:

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

## CLI Integration

### Create CLI for Your Detector

The framework provides automatic CLI generation:

```python
# 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

```bash
# 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
```

## Logging

The framework uses [structlog](https://www.structlog.org/) for structured logging:

```python
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):
```json
{"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:**

```python
# 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](https://github.com/bolin8017/islab-malware-detector/tree/main/examples) 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](LICENSE.txt) for details.

## Citation

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

```bibtex
@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}
}
```