From f4238127154f5043e8cba975ddf25e83ef4937df Mon Sep 17 00:00:00 2001
From: Julien Loizelet <julienloizelet@okaeli.com>
Date: Thu, 8 Feb 2024 15:06:24 +0900
Subject: [PATCH] feat(prune): Add CAPIClient::prune_failing_machines_signals
 method for deleting signals from failing machines

---
 CHANGELOG.md                               | 17 ++++++-
 examples/prune_failing_machines_signals.py | 58 ++++++++++++++++++++++
 src/cscapi/client.py                       | 28 +++++++----
 3 files changed, 91 insertions(+), 12 deletions(-)
 create mode 100644 examples/prune_failing_machines_signals.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7925752..af659b2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,7 +10,22 @@ functions provided by the `src/cscapi` folder.
 
 --- 
 
-## [0.0.2](https://github.com/crowdsecurity/python-capi-sdk/releases/tag/v1.1.0) - 2024-02-07
+
+## [0.1.0](https://github.com/crowdsecurity/python-capi-sdk/releases/tag/v0.1.0) - 2024-02-??
+[_Compare with previous release_](https://github.com/crowdsecurity/python-capi-sdk/compare/v0.0.2...v0.1.0)
+
+### Changed
+
+- **Breaking change**: Change method name `CAPIClient::has_valid_scenarios` to `CAPIClient::_has_valid_scenarios`
+
+### Added
+
+- Add `CAPIClient::prune_failing_machines_signals` method for deleting signals from failing machines
+
+
+---
+
+## [0.0.2](https://github.com/crowdsecurity/python-capi-sdk/releases/tag/v0.0.2) - 2024-02-07
 [_Compare with previous release_](https://github.com/crowdsecurity/python-capi-sdk/compare/v0.0.1...v0.0.2)
 
 
diff --git a/examples/prune_failing_machines_signals.py b/examples/prune_failing_machines_signals.py
new file mode 100644
index 0000000..a39bd6f
--- /dev/null
+++ b/examples/prune_failing_machines_signals.py
@@ -0,0 +1,58 @@
+"""
+This script deletes signals linked to a failing machine.
+"""
+
+import argparse
+import sys
+
+from cscapi.client import CAPIClient, CAPIClientConfig
+from cscapi.sql_storage import SQLStorage
+
+
+class CustomHelpFormatter(argparse.HelpFormatter):
+    def __init__(self, prog, indent_increment=2, max_help_position=48, width=None):
+        super().__init__(prog, indent_increment, max_help_position, width)
+
+
+parser = argparse.ArgumentParser(
+    description="Script to prune failing machines signals.",
+    formatter_class=CustomHelpFormatter,
+)
+
+try:
+    parser.add_argument(
+        "--database",
+        type=str,
+        help="Local database name. Example: cscapi.db",
+        required=True,
+    )
+    args = parser.parse_args()
+except argparse.ArgumentError as e:
+    print(e)
+    parser.print_usage()
+    sys.exit(2)
+
+database = args.database
+database_message = f"\tLocal storage database: {database}\n"
+
+print(
+    f"\nPruning signals for failing machines\n\n"
+    f"Details:\n"
+    f"{database_message}"
+    f"\n\n"
+)
+
+confirmation = input("Do you want to proceed? (Y/n): ")
+if confirmation.lower() == "n":
+    print("Operation cancelled by the user.")
+    sys.exit()
+
+client = CAPIClient(
+    storage=SQLStorage(connection_string=f"sqlite:///{database}"),
+    config=CAPIClientConfig(
+        scenarios=[],
+    ),
+)
+
+
+client.prune_failing_machines_signals()
diff --git a/src/cscapi/client.py b/src/cscapi/client.py
index 7cf734e..921726b 100644
--- a/src/cscapi/client.py
+++ b/src/cscapi/client.py
@@ -9,9 +9,8 @@
 
 import httpx
 import jwt
-from more_itertools import batched
-
 from cscapi.storage import MachineModel, ReceivedDecision, SignalModel, StorageInterface
+from more_itertools import batched
 
 __version__ = metadata.version("cscapi").split("+")[0]
 
@@ -78,24 +77,31 @@ def __init__(self, storage: StorageInterface, config: CAPIClientConfig):
             {"User-Agent": f"{config.user_agent_prefix}-capi-py-sdk/{__version__}"}
         )
 
-    def has_valid_scenarios(self, machine: MachineModel) -> bool:
-        current_scenarios = self.scenarios
-        stored_scenarios = machine.scenarios
-        if len(stored_scenarios) == 0:
-            return False
-
-        return current_scenarios == stored_scenarios
-
     def add_signals(self, signals: List[SignalModel]):
         for signal in signals:
             self.storage.update_or_create_signal(signal)
 
+    def prune_failing_machines_signals(self):
+        signals = self.storage.get_all_signals()
+        for machine_id, signals in _group_signals_by_machine_id(signals).items():
+            machine = self.storage.get_machine_by_id(machine_id)
+            if machine.is_failing:
+                self.storage.delete_signals(signals)
+
     def send_signals(self, prune_after_send: bool = True):
         unsent_signals_by_machineid = _group_signals_by_machine_id(
             filter(lambda signal: not signal.sent, self.storage.get_all_signals())
         )
         self._send_signals_by_machine_id(unsent_signals_by_machineid, prune_after_send)
 
+    def _has_valid_scenarios(self, machine: MachineModel) -> bool:
+        current_scenarios = self.scenarios
+        stored_scenarios = machine.scenarios
+        if len(stored_scenarios) == 0:
+            return False
+
+        return current_scenarios == stored_scenarios
+
     def _send_signals_by_machine_id(
         self,
         signals_by_machineid: Dict[str, List[SignalModel]],
@@ -287,7 +293,7 @@ def _ensure_machine_capi_registered(self, machine: MachineModel) -> MachineModel
     def _ensure_machine_capi_connected(self, machine: MachineModel) -> MachineModel:
         if not has_valid_token(
             machine, self.latency_offset
-        ) or not self.has_valid_scenarios(machine):
+        ) or not self._has_valid_scenarios(machine):
             return self._refresh_machine_token(machine)
         return machine