diff --git a/README.md b/README.md index 89941ac..7d54fb2 100644 --- a/README.md +++ b/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//`. +- 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). diff --git a/bin/aip b/bin/aip index d28055b..f306df5 100755 --- a/bin/aip +++ b/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() diff --git a/etc/docker/README.md b/etc/docker/README.md index 88818af..5e4fccf 100644 --- a/etc/docker/README.md +++ b/etc/docker/README.md @@ -1,40 +1,51 @@ # 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'] @@ -42,4 +53,6 @@ An example output is shown below: 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'] ... -``` \ No newline at end of file +``` + +Without `-v` or `-d`, AIP only prints errors. Blocklists are saved in `data/output/`. See the [main README](../../README.md#commands) for all commands. diff --git a/lib/aip/models/alpha.py b/lib/aip/models/alpha.py index 5ae32b4..44ce497 100644 --- a/lib/aip/models/alpha.py +++ b/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.") diff --git a/lib/aip/models/alpha7.py b/lib/aip/models/alpha7.py new file mode 100644 index 0000000..fc2b918 --- /dev/null +++ b/lib/aip/models/alpha7.py @@ -0,0 +1,34 @@ +""" +AIP - AlphaX Model + +This model filter the attackers of the previous X days. + +This program is free software: you can redistribute it and/or modify it under +the terms of the GNU General Public License as published by the Free Software +Foundation, either version 3 of the License, or (at your option) any later +version. + +This program is distributed in the hope that it will be useful, but WITHOUT +ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. +You should have received a copy of the GNU General Public License along with +this program. If not, see . +""" + +__authors__ = ["Joaquin Bogado "] +__contact__ = "stratosphere@aic.fel.cvut.cz" +__copyright__ = "Copyright 2022, Stratosphere Laboratory." +__credits__ = ["Joaquín Bogado"] +__deprecated__ = False +__license__ = "GPLv3" +__maintainer__ = "Joaquin Bogado" +__version__ = "0.0.1" + +from aip.models.alpha import Alpha +from aip.utils.autoload import register + +@register +class Alpha7(Alpha): + def __init__(self, lookback=7): + super().__init__() + self.lookback = lookback diff --git a/lib/aip/models/base.py b/lib/aip/models/base.py index 4539cac..40c884a 100644 --- a/lib/aip/models/base.py +++ b/lib/aip/models/base.py @@ -26,10 +26,11 @@ import pandas as pd import logging +from os import path from aip.data.access import data_dir -from aip.utils.autoload import register, models -from os import path +from aip.utils.autoload import models +from aip.utils.autoload import register # @register class BaseModel(): diff --git a/lib/aip/models/prioritize.py b/lib/aip/models/prioritize.py deleted file mode 100644 index e666a22..0000000 --- a/lib/aip/models/prioritize.py +++ /dev/null @@ -1,266 +0,0 @@ -""" -AIP - Prioritize New/Consistent Algorithm - -This module prioritize the new/consistent attackers. - -This program is free software: you can redistribute it and/or modify it under -the terms of the GNU General Public License as published by the Free Software -Foundation, either version 3 of the License, or (at your option) any later -version. - -This program is distributed in the hope that it will be useful, but WITHOUT -ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS -FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. -You should have received a copy of the GNU General Public License along with -this program. If not, see . -""" - -__authors__ = ["Joaquin Bogado "] -__contact__ = "stratosphere@aic.fel.cvut.cz" -__copyright__ = "Copyright 2022, Stratosphere Laboratory." -__credits__ = ["Joaquín Bogado"] -__deprecated__ = False -__license__ = "GPLv3" -__maintainer__ = "Joaquin Bogado" -__version__ = "0.0.1" - - -import numpy as np -import pandas as pd -import time - -from aip.data.access import data_path, get_attacks -from aip.models.base import BaseModel -from aip.utils.autoload import register, models -from datetime import date, datetime, timedelta -from os import path -from sklearn.ensemble import RandomForestClassifier - - -def _add_knowledge(last_knowledge, day): - print(f'DEBUG: PROCESSING DATE {day}') - st_time = time.time() - day = datetime.strptime(day, '%Y-%m-%d').date() - p = path.join(data_path, 'processed', 'prioritizers') - attacks = get_attacks(dates=[day]) - df = attacks[0] - if df.empty is True: - last_knowledge.to_csv(path.join(p, f'knowledgebase-{day}-snapshot.gz'), - columns=['orig', 'flows', 'duration','bytes', 'packets', - 'mean_flows', 'mean_duration', 'mean_bytes', 'mean_packets', - 'days_active', 'first_seen', 'last_seen'], - index=False, compression='gzip') - return last_knowledge - df = df.rename(columns={"count": "flows"}) - df.loc[:, 'first_seen'] = day - df.loc[:, 'last_seen'] = day - df.loc[:, 'days_active'] = 1 - df.loc[df['flows'] == 0, 'flows'] = 1 - last_knowledge = pd.concat([last_knowledge, df]) - dates_min = last_knowledge.groupby('orig').first_seen.min() - dates_max = last_knowledge.groupby('orig').last_seen.max() - knowledge = last_knowledge.groupby('orig').sum(numeric_only=True) - knowledge.loc[:,'first_seen'] = dates_min - knowledge.loc[:,'last_seen'] = dates_max - knowledge.loc[:,'mean_flows'] = knowledge['flows']/knowledge['days_active'] - knowledge.loc[:,'mean_duration'] = knowledge['duration']/knowledge['flows'] - knowledge.loc[:,'mean_bytes'] = knowledge['bytes']/knowledge['flows'] - knowledge.loc[:,'mean_packets'] = knowledge['packets']/knowledge['flows'] - knowledge.reset_index(inplace=True) - knowledge.to_csv(path.join(p, f'knowledgebase-{day}-snapshot.gz'), - columns=['orig', 'flows', 'duration','bytes', 'packets', - 'mean_flows', 'mean_duration', 'mean_bytes', 'mean_packets', - 'days_active', 'first_seen', 'last_seen'], - index=False, compression='gzip') - print(f'DEBUG: PROCESSED IN {(time.time() - st_time)/60} MINUTES.') - return knowledge - -def _build_knowledge(start, end): - dates = pd.date_range(start=start, end=end) - last_knowledge = pd.DataFrame() - for day in dates: - last_knowledge = _add_knowledge(last_knowledge, str(day.date())) - return last_knowledge - - -class Knowledgebase(): - ''' - Object that represents the IP features database - ''' - - def _check_date_param(self, day_unchk): - if day_unchk == 'yesterday': - day = str(date.today() - timedelta(days=1)) - elif type(day_unchk) is str: - # force is a date or throw and exception - day = str(datetime.strptime(day_unchk, '%Y-%m-%d').date()) - elif type(day_unchk) is date: - day = str(day_unchk) - return day - - def _load_knowledge_until(self, day): - day = self._check_date_param(day) - self.path = path.join(data_path, - 'processed', 'prioritizers', f'knowledgebase-{day}-snapshot.gz') - if not path.exists(self.path): - self.build(end=datetime.strptime(day, '%Y-%m-%d').date()) - self.knowledge = pd.read_csv(self.path) - self.knowledge.loc[:, 'last_seen'] = pd.to_datetime(self.knowledge.last_seen) - self.knowledge.loc[:, 'first_seen'] = pd.to_datetime(self.knowledge.first_seen) - self.timeframe = (self.knowledge.last_seen.min(), self.knowledge.last_seen.max()) - - def __init__(self, load_until='yesterday'): - day = self._check_date_param(load_until) - self.path = path.join(data_path, - 'processed', 'prioritizers', f'knowledgebase-{day}-snapshot.gz') - self._load_knowledge_until(day) - - - def build(self, start=date.today() - timedelta(days=2), end=date.today() - timedelta(days=1), force=False): - if path.exists(self.path) and not force: - print('Knowledge exists already. Use force=True to rebuild it') - return - # check if the snapshot for the start date exists - # if not, all the snapshots must be created again - p = path.join(data_path, 'processed', 'prioritizers') - if not path.exists(path.join(p, f'knowledgebase-{str(start)}-snapshot.gz')): - last_knowledge = _build_knowledge(start=start, end=end) - else: - days_ago = 1 - day = str(end - timedelta(days=days_ago)) - while not path.exists(path.join(p, f'knowledgebase-{day}-snapshot.gz')): - days_ago += 1 - day = str(end - timedelta(days=days_ago)) - last_knowledge = pd.read_csv(path.join(p, f'knowledgebase-{day}-snapshot.gz')) - last_knowledge.loc[:, 'first_seen'] = pd.to_datetime(last_knowledge.first_seen).dt.date - last_knowledge.loc[:, 'last_seen'] = pd.to_datetime(last_knowledge.last_seen).dt.date - while days_ago >= 1: - days_ago -= 1 - day = str(end - timedelta(days=days_ago)) - last_knowledge = _add_knowledge(last_knowledge, day) - - -@register -class Consistent(BaseModel): - ''' - Prioritize Consistent algorithm - ''' - - def __init__(self): - super().__init__() - # Weights from Thomas' Thesis - # nflows, conns, nbytes, npackets, - # mean_flows, mean_conns, mean_bytes, mean_packets - self.weights = [0.05, 0.05, 0.05, 0.05, 0.2, 0.2, 0.2, 0.2] - self.score_threshold = 0.00002 - self.min_ip_number = 5000 - - def _get_IP_scores(self): - days_since_last_seen = np.array([x.days for x in (pd.to_datetime(date.today()) - self.db.knowledge.last_seen)]) - features = self.db.knowledge[ - ['flows', 'duration','bytes', 'packets', 'mean_flows', - 'mean_duration', 'mean_bytes', 'mean_packets'] - ].values - normalized_features = np.zeros_like(features) - for i in range(8): - normalized_features[:,i] = (features[:,i] - features[:,i].min()) / (features[:,i].max() - features[:,i].min()) - ipscores = normalized_features * self.weights - ipscores = ipscores.sum(axis=1) - days_since_last_seen = np.array([x.days for x in (pd.to_datetime(date.today()) - self.db.knowledge.last_seen)]) - total_attack_time = np.array([x.days for x in (self.db.knowledge.last_seen - self.db.knowledge.first_seen)]) - aging = 1 - (days_since_last_seen/(days_since_last_seen + total_attack_time)) - ipscores *= aging - return ipscores - - def run(self, for_date=date.today()): - self.db = Knowledgebase(load_until=for_date - timedelta(days=1)) - ipscores = self._get_IP_scores() - df = pd.DataFrame() - df['ip'] = self.db.knowledge['orig'].values - df['score'] = ipscores - df = df.sort_values(by='score', ascending=False) - df = self.sanitize(df) - if len(df[df.score > self.score_threshold]) < self.min_ip_number: - df = df.iloc[:self.min_ip_number] - else: - df = df[df.score > self.score_threshold] - return df - - -@register -class New(Consistent): - ''' - Prioritize New algorithm - ''' - - def __init__(self): - super().__init__() - # Weights from Thomas' Thesis - # nflows, conns, nbytes, npackets, - # mean_flows, mean_conns, mean_bytes, mean_packets - self.weights = [0.2, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.05] - self.score_threshold = 0.00001 - self.min_ip_number = 5000 - - def _get_IP_scores(self): - days_since_last_seen = np.array([x.days for x in (pd.to_datetime(date.today()) - self.db.knowledge.last_seen)]) - features = self.db.knowledge[ - ['flows', 'duration','bytes', 'packets', 'mean_flows', - 'mean_duration', 'mean_bytes', 'mean_packets'] - ].values - normalized_features = np.zeros_like(features) - for i in range(8): - normalized_features[:,i] = (features[:,i] - features[:,i].min()) / (features[:,i].max() - features[:,i].min()) - ipscores = normalized_features * self.weights - ipscores = ipscores.sum(axis=1) - aging = 2 / (2 + days_since_last_seen) - ipscores *= aging - return ipscores - -class RandomForest(BaseModel): - ''' - Prioritize Random Forest from Thomas O'Hara's thesis - ''' - - def __init__(self): - super().__init__() - - def _get_training_target_set(self, for_date=date.today()): - - target = Knowledgebase(load_until=for_date - timedelta(days=1)) - target.attacks = get_attacks(dates=[(for_date - timedelta(days=1))])[-1] - days_since_last_seen = np.array([x.days for x in (pd.to_datetime(for_date) - target.knowledge.last_seen)]) - total_attack_time = np.array([x.days for x in (target.knowledge.last_seen - target.knowledge.first_seen)]) - target.knowledge['last_attack_days'] = days_since_last_seen - target.knowledge['first_attack_days'] = total_attack_time - target.features = target.knowledge[target.knowledge['orig'].isin(target.attacks['orig'])] - target.features = target.knowledge - target.X = target.features.drop(columns=['orig', 'first_seen', 'last_seen']).values - - training = Knowledgebase(load_until=for_date - timedelta(days=2)) - training.knowledge.loc[training.knowledge['flows'] == 0, 'flows'] = 1 - training.attacks = pd.concat(get_attacks(dates=[(for_date - timedelta(days=2))])) - days_since_last_seen = np.array([x.days for x in (pd.to_datetime(for_date) - training.knowledge.last_seen)]) - total_attack_time = np.array([x.days for x in (training.knowledge.last_seen - training.knowledge.first_seen)]) - training.knowledge['last_attack_days'] = days_since_last_seen - training.knowledge['first_attack_days'] = total_attack_time - training.knowledge['target'] = training.knowledge['orig'].isin(target.attacks['orig']).values.astype(int) - training.features = training.knowledge[training.knowledge['orig'].isin(training.attacks['orig'])] - training.features = training.knowledge - training.X = training.features.drop(columns=['orig', 'first_seen', 'last_seen', 'target']).values - training.Y = training.features['target'].values - return training, target - - def run(self, for_date=date.today()): - training, target = self._get_training_target_set(for_date=for_date) - clf = RandomForestClassifier(n_estimators=30) - if (target.X.shape[0] > 0) and (training.X.shape[0] > 0): - clf = clf.fit(training.X, training.Y) - pred = clf.predict(target.X) - df = target.features[pred.astype(bool)] - df = df.rename(columns={'orig': 'ip'}) - df = self.sanitize(df) - return df[['ip']] - else: - df = pd.DataFrame(columns=[['ip']]) - return df diff --git a/lib/aip/models/prioritize_consistent.py b/lib/aip/models/prioritize_consistent.py new file mode 100644 index 0000000..b465aa6 --- /dev/null +++ b/lib/aip/models/prioritize_consistent.py @@ -0,0 +1,82 @@ +""" +AIP - Prioritize New/Consistent Algorithm + +This module prioritize the new/consistent attackers. + +This program is free software: you can redistribute it and/or modify it under +the terms of the GNU General Public License as published by the Free Software +Foundation, either version 3 of the License, or (at your option) any later +version. + +This program is distributed in the hope that it will be useful, but WITHOUT +ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. +You should have received a copy of the GNU General Public License along with +this program. If not, see . +""" + +__authors__ = ["Joaquin Bogado "] +__contact__ = "stratosphere@aic.fel.cvut.cz" +__copyright__ = "Copyright 2022, Stratosphere Laboratory." +__credits__ = ["Joaquín Bogado"] +__deprecated__ = False +__license__ = "GPLv3" +__maintainer__ = "Joaquin Bogado" +__version__ = "0.0.1" + + +import numpy as np +import pandas as pd + +from aip.models.base import BaseModel +from aip.utils.autoload import register +from aip.utils.knowledge_base import Knowledgebase +from datetime import date +from datetime import timedelta + + +@register +class Consistent(BaseModel): + ''' + Prioritize Consistent algorithm + ''' + + def __init__(self): + super().__init__() + # Weights from Thomas' Thesis + # nflows, conns, nbytes, npackets, + # mean_flows, mean_conns, mean_bytes, mean_packets + self.weights = [0.05, 0.05, 0.05, 0.05, 0.2, 0.2, 0.2, 0.2] + self.score_threshold = 0.00002 + self.min_ip_number = 5000 + + def _get_IP_scores(self): + days_since_last_seen = np.array([x.days for x in (pd.to_datetime(date.today()) - self.db.knowledge.last_seen)]) + features = self.db.knowledge[ + ['flows', 'duration','bytes', 'packets', 'mean_flows', + 'mean_duration', 'mean_bytes', 'mean_packets'] + ].values + normalized_features = np.zeros_like(features) + for i in range(8): + normalized_features[:,i] = (features[:,i] - features[:,i].min()) / (features[:,i].max() - features[:,i].min()) + ipscores = normalized_features * self.weights + ipscores = ipscores.sum(axis=1) + days_since_last_seen = np.array([x.days for x in (pd.to_datetime(date.today()) - self.db.knowledge.last_seen)]) + total_attack_time = np.array([x.days for x in (self.db.knowledge.last_seen - self.db.knowledge.first_seen)]) + aging = 1 - (days_since_last_seen/(days_since_last_seen + total_attack_time)) + ipscores *= aging + return ipscores + + def run(self, for_date=date.today()): + self.db = Knowledgebase(load_until=for_date - timedelta(days=1)) + ipscores = self._get_IP_scores() + df = pd.DataFrame() + df['ip'] = self.db.knowledge['orig'].values + df['score'] = ipscores + df = df.sort_values(by='score', ascending=False) + df = self.sanitize(df) + if len(df[df.score > self.score_threshold]) < self.min_ip_number: + df = df.iloc[:self.min_ip_number] + else: + df = df[df.score > self.score_threshold] + return df diff --git a/lib/aip/models/prioritize_new.py b/lib/aip/models/prioritize_new.py new file mode 100644 index 0000000..f31be18 --- /dev/null +++ b/lib/aip/models/prioritize_new.py @@ -0,0 +1,64 @@ +""" +AIP - Prioritize New/Consistent Algorithm + +This module prioritize the new/consistent attackers. + +This program is free software: you can redistribute it and/or modify it under +the terms of the GNU General Public License as published by the Free Software +Foundation, either version 3 of the License, or (at your option) any later +version. + +This program is distributed in the hope that it will be useful, but WITHOUT +ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. +You should have received a copy of the GNU General Public License along with +this program. If not, see . +""" + +__authors__ = ["Joaquin Bogado "] +__contact__ = "stratosphere@aic.fel.cvut.cz" +__copyright__ = "Copyright 2022, Stratosphere Laboratory." +__credits__ = ["Joaquín Bogado"] +__deprecated__ = False +__license__ = "GPLv3" +__maintainer__ = "Joaquin Bogado" +__version__ = "0.0.1" + + +import numpy as np +import pandas as pd + +from aip.models.prioritize_consistent import Consistent +from aip.utils.autoload import register +from datetime import date + + +@register +class New(Consistent): + ''' + Prioritize New algorithm + ''' + + def __init__(self): + super().__init__() + # Weights from Thomas' Thesis + # nflows, conns, nbytes, npackets, + # mean_flows, mean_conns, mean_bytes, mean_packets + self.weights = [0.2, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.05] + self.score_threshold = 0.00001 + self.min_ip_number = 5000 + + def _get_IP_scores(self): + days_since_last_seen = np.array([x.days for x in (pd.to_datetime(date.today()) - self.db.knowledge.last_seen)]) + features = self.db.knowledge[ + ['flows', 'duration','bytes', 'packets', 'mean_flows', + 'mean_duration', 'mean_bytes', 'mean_packets'] + ].values + normalized_features = np.zeros_like(features) + for i in range(8): + normalized_features[:,i] = (features[:,i] - features[:,i].min()) / (features[:,i].max() - features[:,i].min()) + ipscores = normalized_features * self.weights + ipscores = ipscores.sum(axis=1) + aging = 2 / (2 + days_since_last_seen) + ipscores *= aging + return ipscores diff --git a/lib/aip/models/random_forest.py b/lib/aip/models/random_forest.py new file mode 100644 index 0000000..64e448a --- /dev/null +++ b/lib/aip/models/random_forest.py @@ -0,0 +1,85 @@ +""" +AIP - Prioritize New/Consistent Algorithm + +This module prioritize the new/consistent attackers. + +This program is free software: you can redistribute it and/or modify it under +the terms of the GNU General Public License as published by the Free Software +Foundation, either version 3 of the License, or (at your option) any later +version. + +This program is distributed in the hope that it will be useful, but WITHOUT +ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. +You should have received a copy of the GNU General Public License along with +this program. If not, see . +""" + +__authors__ = ["Joaquin Bogado "] +__contact__ = "stratosphere@aic.fel.cvut.cz" +__copyright__ = "Copyright 2022, Stratosphere Laboratory." +__credits__ = ["Joaquín Bogado"] +__deprecated__ = False +__license__ = "GPLv3" +__maintainer__ = "Joaquin Bogado" +__version__ = "0.0.1" + + +import numpy as np +import pandas as pd + +from aip.data.access import get_attacks +from aip.models.base import BaseModel +from aip.utils.knowledge_base import Knowledgebase +from datetime import date +from datetime import timedelta +from sklearn.ensemble import RandomForestClassifier + + +class RandomForest(BaseModel): + ''' + Prioritize Random Forest from Thomas O'Hara's thesis + ''' + + def __init__(self): + super().__init__() + + def _get_training_target_set(self, for_date=date.today()): + + target = Knowledgebase(load_until=for_date - timedelta(days=1)) + target.attacks = get_attacks(dates=[(for_date - timedelta(days=1))])[-1] + days_since_last_seen = np.array([x.days for x in (pd.to_datetime(for_date) - target.knowledge.last_seen)]) + total_attack_time = np.array([x.days for x in (target.knowledge.last_seen - target.knowledge.first_seen)]) + target.knowledge['last_attack_days'] = days_since_last_seen + target.knowledge['first_attack_days'] = total_attack_time + target.features = target.knowledge[target.knowledge['orig'].isin(target.attacks['orig'])] + target.features = target.knowledge + target.X = target.features.drop(columns=['orig', 'first_seen', 'last_seen']).values + + training = Knowledgebase(load_until=for_date - timedelta(days=2)) + training.knowledge.loc[training.knowledge['flows'] == 0, 'flows'] = 1 + training.attacks = pd.concat(get_attacks(dates=[(for_date - timedelta(days=2))])) + days_since_last_seen = np.array([x.days for x in (pd.to_datetime(for_date) - training.knowledge.last_seen)]) + total_attack_time = np.array([x.days for x in (training.knowledge.last_seen - training.knowledge.first_seen)]) + training.knowledge['last_attack_days'] = days_since_last_seen + training.knowledge['first_attack_days'] = total_attack_time + training.knowledge['target'] = training.knowledge['orig'].isin(target.attacks['orig']).values.astype(int) + training.features = training.knowledge[training.knowledge['orig'].isin(training.attacks['orig'])] + training.features = training.knowledge + training.X = training.features.drop(columns=['orig', 'first_seen', 'last_seen', 'target']).values + training.Y = training.features['target'].values + return training, target + + def run(self, for_date=date.today()): + training, target = self._get_training_target_set(for_date=for_date) + clf = RandomForestClassifier(n_estimators=30) + if (target.X.shape[0] > 0) and (training.X.shape[0] > 0): + clf = clf.fit(training.X, training.Y) + pred = clf.predict(target.X) + df = target.features[pred.astype(bool)] + df = df.rename(columns={'orig': 'ip'}) + df = self.sanitize(df) + return df[['ip']] + else: + df = pd.DataFrame(columns=[['ip']]) + return df diff --git a/lib/aip/utils/date_utils.py b/lib/aip/utils/date_utils.py index 6cdf829..ea6d408 100644 --- a/lib/aip/utils/date_utils.py +++ b/lib/aip/utils/date_utils.py @@ -1,5 +1,5 @@ -import logging from datetime import datetime +from datetime import date def validate_and_convert_date(date_str): @@ -7,8 +7,9 @@ def validate_and_convert_date(date_str): Validates a date string in 'YYYY-MM-DD' format and converts it to a date object. """ try: - dateobj = datetime.strptime(date_str, '%Y-%m-%d') - return dateobj.date() - except ValueError as e: - logging.error(f"Invalid date format for '{date_str}', expected YYYY-MM-DD") - raise ValueError(f"Invalid date format: {date_str}, expected YYYY-MM-DD") from e \ No newline at end of file + dateobj = datetime.strptime(date_str, '%Y-%m-%d').date() + except ValueError as err: + raise ValueError(f"Invalid date format: {date_str}, expected YYYY-MM-DD") from err + if dateobj > date.today(): + raise ValueError(f"The input date '{date_str}' cannot be in the future.") + return dateobj diff --git a/lib/aip/utils/generate_historical_blocklists.py b/lib/aip/utils/generate_historical_blocklists.py index e56908a..f55debf 100755 --- a/lib/aip/utils/generate_historical_blocklists.py +++ b/lib/aip/utils/generate_historical_blocklists.py @@ -29,10 +29,10 @@ from aip.data.access import data_path, project_dir 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.prioritize import Knowledgebase +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.models.pareto import Pareto from datetime import date, timedelta from joblib import Parallel, delayed diff --git a/lib/aip/utils/knowledge_base.py b/lib/aip/utils/knowledge_base.py new file mode 100644 index 0000000..754b150 --- /dev/null +++ b/lib/aip/utils/knowledge_base.py @@ -0,0 +1,161 @@ +""" +This module provides functionality for managing the +knowledge base used by AIP models. +""" + +import logging +import time +import pandas as pd +from datetime import date +from datetime import datetime +from datetime import timedelta +from os import path +from aip.data.access import data_path +from aip.data.access import get_attacks + + +def _add_knowledge(last_knowledge, day, log_level=logging.ERROR): + # Configure logger + logger = logging.getLogger(__name__) + logging.basicConfig(level=log_level) + + logger.debug(f'_add_knowledge - Processing date {day}') + st_time = time.time() + day = datetime.strptime(day, '%Y-%m-%d').date() + + p = path.join(data_path, 'processed', 'prioritizers') + + attacks = get_attacks(dates=[day]) + + df = attacks[0] + if df.empty is True: + last_knowledge.to_csv(path.join(p, f'knowledgebase-{day}-snapshot.gz'), + columns=['orig', 'flows', 'duration','bytes', 'packets', + 'mean_flows', 'mean_duration', 'mean_bytes', 'mean_packets', + 'days_active', 'first_seen', 'last_seen'], + index=False, compression='gzip') + return last_knowledge + df = df.rename(columns={"count": "flows"}) + df.loc[:, 'first_seen'] = day + df.loc[:, 'last_seen'] = day + df.loc[:, 'days_active'] = 1 + df.loc[df['flows'] == 0, 'flows'] = 1 + + last_knowledge = pd.concat([last_knowledge, df]) + dates_min = last_knowledge.groupby('orig').first_seen.min() + dates_max = last_knowledge.groupby('orig').last_seen.max() + knowledge = last_knowledge.groupby('orig').sum(numeric_only=True) + knowledge.loc[:,'first_seen'] = dates_min + knowledge.loc[:,'last_seen'] = dates_max + knowledge.loc[:,'mean_flows'] = knowledge['flows']/knowledge['days_active'] + knowledge.loc[:,'mean_duration'] = knowledge['duration']/knowledge['flows'] + knowledge.loc[:,'mean_bytes'] = knowledge['bytes']/knowledge['flows'] + knowledge.loc[:,'mean_packets'] = knowledge['packets']/knowledge['flows'] + knowledge.reset_index(inplace=True) + knowledge.to_csv(path.join(p, f'knowledgebase-{day}-snapshot.gz'), + columns=['orig', 'flows', 'duration','bytes', 'packets', + 'mean_flows', 'mean_duration', 'mean_bytes', 'mean_packets', + 'days_active', 'first_seen', 'last_seen'], + index=False, compression='gzip') + + logger.debug(f'_add_knowledge - Processed in {(time.time() - st_time)/60} minutes') + return knowledge + + +def _build_knowledge(start, end, log_level=logging.ERROR): + # Configure logger + logger = logging.getLogger(__name__) + logging.basicConfig(level=log_level) + + dates = pd.date_range(start=start, end=end) + last_knowledge = pd.DataFrame() + for day in dates: + logger.debug(f'_build_knowledge - Building knowledge base for {day}') + last_knowledge = _add_knowledge(last_knowledge, str(day.date()), log_level=log_level) + return last_knowledge + + +def _rebuild(start_date, end_date=None, log_level=logging.ERROR): + """ + Wrapper to rebuild the knowledge base from the specified start date to + yesterday, the last complete day. + """ + # Configure logger + logger = logging.getLogger(__name__) + logging.basicConfig(level=log_level) + + # Today's data is incomplete, do not add it to the knowledge base + if end_date is None: + end_date = date.today() - timedelta(days=1) + + start_time = time.time() + logger.debug(f"_rebuild - Rebuilding knowledge base from {start_date} to {end_date}") + + _build_knowledge(start=start_date, end=end_date, log_level=log_level) + + logger.debug(f'_rebuild - Rebuilding knowledge base finished in {(time.time() - start_time)/60} minutes') + + +class Knowledgebase(): + ''' + Object that represents the IP features database + ''' + + def _check_date_param(self, day_unchk): + if day_unchk == 'yesterday': + day = str(date.today() - timedelta(days=1)) + elif type(day_unchk) is str: + # force is a date or throw and exception + day = str(datetime.strptime(day_unchk, '%Y-%m-%d').date()) + elif type(day_unchk) is date: + day = str(day_unchk) + return day + + def _load_knowledge_until(self, day): + day = self._check_date_param(day) + self.path = path.join(data_path, + 'processed', 'prioritizers', f'knowledgebase-{day}-snapshot.gz') + if not path.exists(self.path): + self.build(end=datetime.strptime(day, '%Y-%m-%d').date()) + self.knowledge = pd.read_csv(self.path) + self.knowledge.loc[:, 'last_seen'] = pd.to_datetime(self.knowledge.last_seen) + self.knowledge.loc[:, 'first_seen'] = pd.to_datetime(self.knowledge.first_seen) + self.timeframe = (self.knowledge.last_seen.min(), self.knowledge.last_seen.max()) + + def __init__(self, load_until='yesterday'): + # Set up the logger for the class + self.logger = logging.getLogger(self.__class__.__name__) + + day = self._check_date_param(load_until) + + self.path = path.join(data_path, 'processed', 'prioritizers', f'knowledgebase-{day}-snapshot.gz') + self._load_knowledge_until(day) + + + def build(self, start=date.today() - timedelta(days=2), end=date.today() - timedelta(days=1), force=False): + if path.exists(self.path) and not force: + self.logger.error('The knowledge already exists. Use force=True to force its rebuild') + return + + # check if the snapshot for the start date exists + # if not, all the snapshots must be created again + p = path.join(data_path, 'processed', 'prioritizers') + + if not path.exists(path.join(p, f'knowledgebase-{str(start)}-snapshot.gz')): + last_knowledge = _build_knowledge(start=start, end=end) + else: + days_ago = 1 + try: + day = str(end - timedelta(days=days_ago)) + while not path.exists(path.join(p, f'knowledgebase-{day}-snapshot.gz')): + days_ago += 1 + day = str(end - timedelta(days=days_ago)) + last_knowledge = pd.read_csv(path.join(p, f'knowledgebase-{day}-snapshot.gz')) + last_knowledge.loc[:, 'first_seen'] = pd.to_datetime(last_knowledge.first_seen).dt.date + last_knowledge.loc[:, 'last_seen'] = pd.to_datetime(last_knowledge.last_seen).dt.date + while days_ago >= 1: + days_ago -= 1 + day = str(end - timedelta(days=days_ago)) + last_knowledge = _add_knowledge(last_knowledge, day) + except OverflowError as err: + raise Exception(f"There is not Knowledge Base for the model to run (Required: {end})") diff --git a/tests/test_lib_aip_utils_date_utils.py b/tests/test_lib_aip_utils_date_utils.py index 79a663c..677e947 100644 --- a/tests/test_lib_aip_utils_date_utils.py +++ b/tests/test_lib_aip_utils_date_utils.py @@ -1,6 +1,7 @@ import unittest import logging from datetime import date +from datetime import timedelta from lib.aip.utils.date_utils import validate_and_convert_date @@ -48,5 +49,13 @@ def test_edge_case(self): # Test leap years self.assertEqual(validate_and_convert_date("2024-02-29"), date(2024, 2, 29)) + def test_future_date(self): + # Test future date raises exception + future_date = (date.today() + timedelta(days=1)).isoformat() + + with self.assertRaisesRegex(ValueError, "cannot be in the future"): + validate_and_convert_date(future_date) + + if __name__ == '__main__': unittest.main()