Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(prune): Add method for deleting signals from failing machines #14

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
58 changes: 58 additions & 0 deletions examples/prune_failing_machines_signals.py
Original file line number Diff line number Diff line change
@@ -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()
28 changes: 17 additions & 11 deletions src/cscapi/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down Expand Up @@ -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]],
Expand Down Expand Up @@ -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

Expand Down