Merge development into main

README.md

@@ -38,25 +38,76 @@ Each AIP model generates its own blocklist based on a specific criteria. The mai
     - A more experimental approach to increase blocklist efficiency.
 
 
-## AIP Docker
+## Setup
 
-The best way to run AIP right now is using [Docker](etc/docker/README.md).
+AIP can run with Docker (recommended) or directly with Python. Either way, it needs two things:
 
-## Usage
+1. **Zeek logs** in `data/raw/`, one folder per day (for example `data/raw/2024-10-22/conn.*`). With Docker you can point AIP to your Zeek logs folder instead of copying them.
+2. **Your honeypot IPs** in `data/external/`. Copy the example files and edit them:
 
-AIP will automatically attempt to run all the models using the available data. Assuming the Zeek data is located in its usual location:
+    ```bash
+    cp data/external/honeypots_public_ips_example.csv data/external/honeypots_public_ips.csv
+    cp data/external/do_not_block_these_ips_example.csv data/external/do_not_block_these_ips.csv
+    ```
+
+    - `honeypots_public_ips.csv` (required): the public IPs of your honeypots. AIP only looks at traffic sent to these IPs.
+    - `do_not_block_these_ips.csv` (optional): IPs that must never appear in a blocklist.
+
+### Run with Docker
+
+From the AIP folder, assuming your Zeek logs are in `/opt/zeek/logs/`:
+
+```bash
+docker run --rm \
+  -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro \
+  -v ${PWD}/data/:/home/aip/AIP/data/:rw \
+  --name aip stratosphereips/aip:latest bin/aip
+```
+
+This runs all models for today. To use any of the [commands](#commands) below, write them after `bin/aip`. For example, to run all models for one day:
 
 ```bash
-:~$ cd AIP
-:~$ docker run --rm -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro -v ${PWD}/data/:/home/aip/AIP/data/:rw --name aip stratosphereips/aip:latest bin/aip
+docker run --rm \
+  -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro \
+  -v ${PWD}/data/:/home/aip/AIP/data/:rw \
+  --name aip stratosphereips/aip:latest bin/aip run-models --start-date 2026-10-12 --end-date 2026-10-12
 ```
 
-To run AIP for a specific day:
+To build the image yourself, see the [Docker guide](etc/docker/README.md).
+
+### Run with Python
+
 ```bash
-:~$ cd AIP
-:~$ docker run --rm -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro -v ${PWD}/data/:/home/aip/AIP/data/:rw --name aip stratosphereips/aip:latest bin/aip YYYY-MM-DD
+git clone https://github.com/stratosphereips/AIP.git ~/AIP
+cd ~/AIP
+python3 -m venv aip-venv
+source aip-venv/bin/activate
+pip install -r requirements.txt
+export PYTHONPATH=$(pwd)/lib
+bin/aip --help
 ```
 
+Run `source aip-venv/bin/activate` and the `export PYTHONPATH` line again each time you open a new terminal.
+
+## Commands
+
+| What you want | Command |
+|---|---|
+| Run all models for today | `bin/aip` |
+| Run all models for one day | `bin/aip run-models --start-date 2026-10-02 --end-date 2026-10-02` |
+| Run all models for several days | `bin/aip run-models --start-date 2026-10-01 --end-date 2026-10-12` |
+| Run only one model | Add `--model` with `Alpha`, `Alpha7`, `Prioritize_New`, `Prioritize_Consistent` or `Random_Forest` |
+| Rebuild the attacker history from a date until yesterday | `bin/aip rebuild-kb --start-date 2026-10-01` |
+
+Things to know:
+
+- If you leave out `--end-date`, it is today. So `--start-date` alone runs every day from that date until today.
+- The blocklist for a day is built from the attacks seen on the days before it.
+- Blocklists are saved in `data/output/<model name>/`.
+- The attacker history is what the Prioritize and Random Forest models learn from. AIP builds it on its own as it runs.
+- To see what AIP is doing, add `-v` (progress) or `-d` (full detail) right after `bin/aip`, for example `bin/aip -v run-models`.
+- See all options with `bin/aip --help`.
+
 ## License
 
 The Stratosphere AIP tool is licensed under [GNU General Public License v3.0](https://github.com/stratosphereips/AIP/blob/main/LICENSE).

bin/aip

@@ -31,18 +31,22 @@ __version__ = "1.0.0"
 
 import argparse
 import logging
+import sys
 from datetime import date
+from datetime import timedelta
 from os import makedirs
 from os import path
 from aip.data.access import data_path
 from aip.models.alpha import Alpha
-from aip.models.prioritize import New
-from aip.models.prioritize import Consistent
-from aip.models.prioritize import RandomForest
+from aip.models.alpha7 import Alpha7
+from aip.models.prioritize_new import New
+from aip.models.prioritize_consistent import Consistent
+from aip.models.random_forest import RandomForest
 from aip.utils.date_utils import validate_and_convert_date
+from aip.utils.knowledge_base import _rebuild
 
 
-def run_model(aip_model_name, aip_model, date_day):
+def run_model(aip_model_name, aip_model, start_date, log_level=logging.ERROR):
     """
     Run a given model with exception handling
     """
@@ -53,48 +57,113 @@ def run_model(aip_model_name, aip_model, date_day):
         makedirs(model_output_dir)
 
     try:
-        blocklist = aip_model.run(date_day)
-        blocklist.to_csv(path.join(model_output_dir, f'AIP-{aip_model_name}-{str(date_day)}.csv.gz'), index=False, compression='gzip')
+        blocklist = aip_model.run(start_date)
+        blocklist.to_csv(path.join(model_output_dir, f'AIP-{aip_model_name}-{str(start_date)}.csv.gz'), index=False, compression='gzip')
         logging.info(f"{aip_model_name} model completed successfully.")
-    except Exception as e:
-        logging.error(f"Error running {aip_model_name} model: {e}", exc_info=True)
+    except Exception as err:
+        logging.error(f"Error running {aip_model_name} model: {err}")
 
 
 def main():
     parser = argparse.ArgumentParser(description='Attacker IP Prioritization (AIP) Tool')
-    parser.add_argument('--date', type=str, help='The date for running the models in YYYY-MM-DD format. Defaults to today.', default=str(date.today()))
-    parser.add_argument('--model', type=str, choices=['Alpha', 'Alpha7', 'Prioritize_New', 'Prioritize_Consistent', 'Random_Forest', 'all'], default='all', help='Select AIP model to run. Defaults to all.')
     parser.add_argument('-d', '--debug', required=False, help="Debugging mode.", action="store_const", dest="log_level", const=logging.DEBUG, default=logging.ERROR,)
     parser.add_argument('-v', '--verbose', required=False, help="Verbose mode", action="store_const", dest="log_level", const=logging.INFO,)
 
+    subparsers = parser.add_subparsers(dest="command", help="commands")
+
+    # Run Models command 
+    run_models_parser = subparsers.add_parser("run-models", help="Run models with optional date range.")
+    run_models_parser.add_argument('--start-date', type=str, help='The start date for running the models (YYYY-MM-DD). Default is today.', default=str(date.today()))
+    run_models_parser.add_argument('--end-date', type=str, help='The end date for running the models (YYYY-MM-DD). Default is today.', default=str(date.today()))
+    run_models_parser.add_argument('--model', type=str, choices=['Alpha', 'Alpha7', 'Prioritize_New', 'Prioritize_Consistent', 'Random_Forest', 'all'], default='all', help='Select AIP model to run. Default is all models.')
+
+    # Rebuild KB command
+    rebuild_kb_parser = subparsers.add_parser("rebuild-kb", help="Rebuild the knowledge base from a specific date.")
+    rebuild_kb_parser.add_argument("--start-date", required=True, type=str, help="The date from which to rebuild the knowledge base in YYYY-MM-DD format.")
+
     args = parser.parse_args()
 
     # Set up logging
+    logger = logging.getLogger('aip')
     log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
     logging.basicConfig(level=args.log_level, format=log_fmt)
 
-    # Validate input date
-    run_date_day = validate_and_convert_date(args.date)
-
-    # Run Alpha Model
-    if args.model in ['Alpha', 'all']:
-        run_model('Alpha', Alpha(), run_date_day)
-
-    # Alpha 7 Model
-    if args.model in ['Alpha7', 'all']:
-        run_model('Alpha7', Alpha(lookback=7), run_date_day)
-
-    # Prioritize New Model
-    if args.model in ['Prioritize_New', 'all']:
-        run_model('Prioritize_New', New(), run_date_day)
-
-    # Prioritize Consistent Model
-    if args.model in ['Prioritize_Consistent', 'all']:
-        run_model('Prioritize_Consistent', Consistent(), run_date_day)
-
-    # Prioritize Random Forest Model
-    if args.model in ['Random_Forest', 'all']:
-        run_model('Random_Forest', RandomForest(), run_date_day)
+    # If not set, default to today()
+    if not hasattr(args, 'start_date'):
+        args.start_date = date.today().strftime('%Y-%m-%d')
+    if not hasattr(args, 'end_date'):
+        args.end_date = date.today().strftime('%Y-%m-%d')
+
+    try: 
+        # Validate once not on every command
+        aip_start_date = validate_and_convert_date(args.start_date)
+        aip_end_date = validate_and_convert_date(args.end_date)
+        logger.debug(f"AIP will run models from start date {aip_start_date} to end date {aip_end_date}")
+
+        # If model is not set, default to today
+        if not hasattr(args,'model'):
+            args.model = 'all'
+
+        # Rebuild knowledge base
+        if args.command == "rebuild-kb":
+            last_complete_day = date.today() - timedelta(days=1)
+            if aip_start_date > last_complete_day:
+                logger.error(f"The start date must be {last_complete_day} or earlier, today is not complete yet.")
+                sys.exit(1)
+            logger.info(f"Rebuilding knowledge base from {aip_start_date} onwards")
+            _rebuild(start_date=aip_start_date, log_level=args.log_level)
+            sys.exit()
+
+        if aip_start_date > aip_end_date:
+            logger.error(f"The start date {aip_start_date} is after the end date {aip_end_date}.")
+            sys.exit(1)
+
+        while aip_start_date <= aip_end_date:
+            logging.info(f"Running {args.model} model(s) for date {aip_start_date}.")
+            # Run Alpha Model
+            if args.model in ['Alpha', 'all']:
+                run_model(aip_model_name='Alpha',
+                        aip_model=Alpha(),
+                        start_date=aip_start_date,
+                        log_level=args.log_level
+                        )
+
+            # Alpha 7 Model
+            if args.model in ['Alpha7', 'all']:
+                run_model(aip_model_name='Alpha7',
+                        aip_model=Alpha7(),
+                        start_date=aip_start_date,
+                        log_level=args.log_level
+                        )
+
+            # Prioritize New Model
+            if args.model in ['Prioritize_New', 'all']:
+                run_model(aip_model_name='Prioritize_New',
+                        aip_model=New(),
+                        start_date=aip_start_date,
+                        log_level=args.log_level
+                        )
+
+            # Prioritize Consistent Model
+            if args.model in ['Prioritize_Consistent', 'all']:
+                run_model(aip_model_name='Prioritize_Consistent',
+                        aip_model=Consistent(),
+                        start_date=aip_start_date,
+                        log_level=args.log_level
+                        )
+
+            # Prioritize Random Forest Model
+            if args.model in ['Random_Forest', 'all']:
+                run_model(aip_model_name='Random_Forest',
+                        aip_model=RandomForest(),
+                        start_date=aip_start_date,
+                        log_level=args.log_level
+                        )
+            aip_start_date += timedelta(days=1)
+
+    except ValueError as err:
+        logger.error(err)
+        sys.exit(1)
 
 if __name__ == '__main__':
     main()

etc/docker/README.md

@@ -1,45 +1,58 @@
 # Docker image for AIP
 
-AIP docker aims to help in development and deployment of AIP algorithms to newcommers. The code of the repository is mounted and available from inside the docker image. The source of the data and the output folder of AIP is also mounted as a data volume inside the data/ folder to decouple from where this data really is in the host machine, easying the AIP deployment.
+The Docker image includes AIP and everything it needs to run. Your Zeek logs and the AIP `data/` folder stay on your machine and are shared with the container when you run it.
 
 ## Build the image
 
-To build the image, you can run the following command.
-
 ```bash
-:~$ git clone https://github.com/stratosphereips/AIP.git
-:~$ cd AIP/
-:~/AIP$ docker build --build-arg uid=1000 --file etc/docker/Dockerfile --tag aip:latest .
+git clone https://github.com/stratosphereips/AIP.git
+cd AIP/
+docker build --build-arg uid=1000 --file etc/docker/Dockerfile --tag aip:latest .
 ```
 
 ## Prepare the data
 
-AIP needs raw network flow data to run. In this case, we assume you have Zeek logs in `/opt/zeek/logs`. 
+AIP reads Zeek logs. These steps assume your Zeek logs are in `/opt/zeek/logs/`.
 
-Additionally, the following two files need to be edited and populated:
-- `data/external/do_not_block_these_ips_example.csv`: you want to add here IPs that should not appear on the AIP blocklists
-- `data/external/honeypots_public_ips_example.csv`: you want to add here the public IP of the honeypot or machine running Zeek
+Copy the two example files and edit the copies:
 
-First copy the files and then edit them:
 ```bash
-:~/AIP$ cp data/external/do_not_block_these_ips_example.csv data/external/do_not_block_these_ips.csv
-:~/AIP$ cp data/external/honeypots_public_ips_example.csv data/external/honeypots_public_ips.csv
+cp data/external/honeypots_public_ips_example.csv data/external/honeypots_public_ips.csv
+cp data/external/do_not_block_these_ips_example.csv data/external/do_not_block_these_ips.csv
 ```
 
+- `honeypots_public_ips.csv` (required): the public IPs of your honeypots. AIP only looks at traffic sent to these IPs.
+- `do_not_block_these_ips.csv` (optional): IPs that must never appear in a blocklist.
+
 ## Run the container
 
-To run the container of that image you can run the following command:
+Run all models for today:
 
 ```bash
-:~/AIP$ docker run --rm -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro -v ${PWD}/data/:/home/aip/AIP/data/:rw --name aip aip:latest bin/aip
+docker run --rm \
+  -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro \
+  -v ${PWD}/data/:/home/aip/AIP/data/:rw \
+  --name aip aip:latest bin/aip
 ```
 
-An example output is shown below:
+Run all models for one day, showing full detail with `-d`:
+
+```bash
+docker run --rm \
+  -v /opt/zeek/logs/:/home/aip/AIP/data/raw:ro \
+  -v ${PWD}/data/:/home/aip/AIP/data/:rw \
+  --name aip aip:latest bin/aip -d run-models --start-date 2024-10-22 --end-date 2024-10-22
+```
+
+The output starts like this:
+
 ```
 2024-10-23 13:02:36,513 - aip.data.access - DEBUG - Creating attacks for dates ['2024-10-22']
 2024-10-23 13:02:36,513 - aip.data.access - DEBUG - Making  dataset from raw data for dates ['2024-10-22']
 2024-10-23 13:02:37,197 - aip.data.access - DEBUG - Writting file: /home/aip/AIP/data/interim/daily.conn.2024-10-22.csv.gz
 2024-10-23 13:02:37,539 - aip.data.access - DEBUG - Writting file: /home/aip/AIP/data/processed/attacks.2024-10-22.csv.gz
 2024-10-23 13:02:37,648 - aip.data.access - DEBUG - Creating attacks for dates ['2024-10-20', '2024-10-21']
 ...
-```
+```
+
+Without `-v` or `-d`, AIP only prints errors. Blocklists are saved in `data/output/`. See the [main README](../../README.md#commands) for all commands.

lib/aip/models/alpha.py

@@ -27,11 +27,11 @@
 import pandas as pd
 
 from aip.models.base import BaseModel
-from aip.utils.autoload import register, models
+from aip.utils.autoload import register
 from aip.data.access import get_attacks
-from datetime import date, timedelta
+from datetime import date
+from datetime import timedelta
 
-#__models__ = []
 
 @register
 class Alpha(BaseModel):
@@ -42,16 +42,18 @@ def __init__(self, lookback=1):
     def run(self, for_date=date.today()):
         start = str(for_date - timedelta(days=self.lookback))
         end = str(for_date - timedelta(days=1))
-        # get all the attackers IPs
+
+        # Get all the attackers IPs
         attacks = get_attacks(start, end, usecols=['orig'])
         attacks = pd.concat(attacks).drop_duplicates()
-        attacks = attacks.rename(columns={'orig':'ip'})
-        self.blocklist = attacks
-        self.sanitize()
-        return self.blocklist
+        if not attacks.empty:
+            attacks = attacks.rename(columns={'orig':'ip'})
 
-@register
-class Alpha7(Alpha):
-    def __init__(self, lookback=7):
-        super().__init__()
-        self.lookback = lookback
+            self.blocklist = attacks
+
+            # Remove IPs from do_not_block_these_ips.csv
+            self.sanitize()
+
+            return self.blocklist
+        else:
+            raise ValueError("Please check data availability and try again.")