Skip to content

Commit

Permalink
feat(enroll): Fix missing scenario for a fresh enrolled machine
Browse files Browse the repository at this point in the history
  • Loading branch information
@julienloizelet
julienloizelet committed Feb 1, 2024
1 parent c2ca90d commit 7955768
Show file tree
Hide file tree
Showing 4 changed files with 77 additions and 26 deletions.
6 changes: 3 additions & 3 deletions examples/basic.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,9 @@
client = CAPIClient(
storage=SQLStorage(),
config=CAPIClientConfig(
scenarios=["pysdktest/test-01", "pysdktest/test-02"],
scenarios=["acme/http-bf", "crowdsec/ssh-bf"],
user_agent_prefix="example",
prod=True,
prod=False,
),
)

Expand All @@ -17,7 +17,7 @@
attacker_ip="81.81.81.81",
scenario="pysdktest/test-sc",
created_at="2024-01-19 12:12:21 +0000",
machine_id=generate_machine_id_from_key("myMachineId"),
machine_id=generate_machine_id_from_key("myMachineKeyIdentifier"),
context=[{"key": "scenario-version", "value": "1.0.0"}],
message="test message to see where it is written",
)
Expand Down
27 changes: 23 additions & 4 deletions examples/enroll.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import sys
from cscapi.client import CAPIClient, CAPIClientConfig
from cscapi.sql_storage import SQLStorage
from cscapi.utils import generate_machine_id_from_key


class CustomHelpFormatter(argparse.HelpFormatter):
Expand All @@ -23,7 +24,10 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
parser.add_argument("--prod", action="store_true", help="Use production mode")
parser.add_argument("--key", type=str, help="Enrollment key to use", required=True)
parser.add_argument(
"--machine_id", type=str, help="ID of the machine", required=True
"--human_machine_id",
type=str,
help="Human readable machine identifier. Will be converted in CrowdSec ID. Example: 'myMachineId'",
required=True,
)
parser.add_argument("--name", type=str, help="Name of the machine", default=None)
parser.add_argument("--overwrite", action="store_true", help="Force overwrite")
Expand All @@ -42,6 +46,12 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
parser.add_argument(
"--user_agent_prefix", type=str, help="User agent prefix", default=None
)
parser.add_argument(
"--database",
type=str,
help="Local database name. Example: cscapi.db",
default=None,
)
args = parser.parse_args()
except argparse.ArgumentError as e:
print(e)
Expand All @@ -50,6 +60,8 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):

tags = json.loads(args.tags) if args.tags else None
scenarios = json.loads(args.scenarios) if args.scenarios else None
machine_id = generate_machine_id_from_key(args.human_machine_id)
machine_id_message = f"\tMachine ID: '{machine_id}'\n"
name_message = f" '{args.name}'" if args.name else ""
user_agent_message = (
f"\tUser agent prefix:'{args.user_agent_prefix}'\n"
Expand All @@ -61,13 +73,20 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
scenarios_message = f"\tScenarios:{args.scenarios}\n" if scenarios else ""
env_message = "\tEnv: production\n" if args.prod else "\tEnv: development\n"

database = "cscapi_examples.db" if args.prod else "cscapi_examples_dev.db"
database = (
args.database
if args.database
else "cscapi_examples_prod.db"
if args.prod
else "cscapi_examples_dev.db"
)
database_message = f"\tLocal storage database: {database}\n"

print(
f"\nEnrolling machine{name_message} with key '{args.key}' and id '{args.machine_id}' {overwrite_message}\n\n"
f"\nEnrolling machine{name_message} with key '{args.key}' {overwrite_message}\n\n"
f"Details:\n"
f"{env_message}"
f"{machine_id_message}"
f"{scenarios_message}"
f"{tags_message}"
f"{user_agent_message}"
Expand All @@ -90,7 +109,7 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
)

client.enroll_machines(
machine_ids=[args.machine_id],
machine_ids=[machine_id],
attachment_key=args.key,
name=args.name,
overwrite=args.overwrite,
Expand Down
27 changes: 22 additions & 5 deletions examples/send_signal.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from cscapi.client import CAPIClient, CAPIClientConfig
from cscapi.sql_storage import SQLStorage
from cscapi.utils import create_signal
from cscapi.utils import generate_machine_id_from_key


class CustomHelpFormatter(argparse.HelpFormatter):
Expand All @@ -23,7 +24,10 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
try:
parser.add_argument("--prod", action="store_true", help="Use production mode")
parser.add_argument(
"--machine_id", type=str, help="ID of the machine", required=True
"--human_machine_id",
type=str,
help="Human readable machine identifier. Will be converted in CrowdSec ID. Example: 'myMachineId'",
required=True,
)
parser.add_argument("--ip", type=str, help="Attacker IP", required=True)
parser.add_argument(
Expand All @@ -47,12 +51,19 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
parser.add_argument(
"--user_agent_prefix", type=str, help="User agent prefix", default=None
)
parser.add_argument(
"--database",
type=str,
help="Local database name. Example: cscapi.db",
default=None,
)
args = parser.parse_args()
except argparse.ArgumentError as e:
print(e)
parser.print_usage()
sys.exit(2)

machine_id = generate_machine_id_from_key(args.human_machine_id)
machine_id_message = f"machine ID: '{machine_id}'\n"
ip_message = f"\tAttacker IP: '{args.ip}'\n"
created_at_message = f"\tCreated at: '{args.created_at}'\n"
scenario_message = f"\tScenario: '{args.scenario}'\n"
Expand All @@ -69,11 +80,17 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
)
env_message = "\tEnv: production\n" if args.prod else "\tEnv: development\n"

database = "cscapi_examples.db" if args.prod else "cscapi_examples_dev.db"
database = (
args.database
if args.database
else "cscapi_examples_prod.db"
if args.prod
else "cscapi_examples_dev.db"
)
database_message = f"\tLocal storage database: {database}\n"

print(
f"\nSending signal for machine '{args.machine_id}'\n\n"
f"\nSending signal for '{machine_id_message}'\n\n"
f"Details:\n"
f"{env_message}"
f"{ip_message}"
Expand Down Expand Up @@ -104,7 +121,7 @@ def __init__(self, prog, indent_increment=2, max_help_position=36, width=None):
attacker_ip=args.ip,
scenario=args.scenario,
created_at=args.created_at,
machine_id=args.machine_id,
machine_id=machine_id,
)
]

Expand Down
43 changes: 29 additions & 14 deletions src/cscapi/client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,15 @@ class CAPIClientConfig:
force_config_scenario: bool = True


def _group_signals_by_machine_id(
signals: Iterable[SignalModel],
) -> Dict[str, List[SignalModel]]:
signals_by_machineid: Dict[str, List[SignalModel]] = defaultdict(list)
for signal in signals:
signals_by_machineid[signal.machine_id].append(signal)
return signals_by_machineid


class CAPIClient:
def __init__(self, storage: StorageInterface, config: CAPIClientConfig):
self.storage = storage
Expand All @@ -71,24 +80,27 @@ 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

if not self.force_config_scenario:
return True

return current_scenarios == stored_scenarios

def add_signals(self, signals: List[SignalModel]):
for signal in signals:
self.storage.update_or_create_signal(signal)

def send_signals(self, prune_after_send: bool = True):
unsent_signals_by_machineid = self._group_signals_by_machine_id(
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 _group_signals_by_machine_id(
self, signals: Iterable[SignalModel]
) -> Dict[str, List[SignalModel]]:
signals_by_machineid: Dict[str, List[SignalModel]] = defaultdict(list)
for signal in signals:
signals_by_machineid[signal.machine_id].append(signal)
return signals_by_machineid

def _send_signals_by_machine_id(
self,
signals_by_machineid: Dict[str, List[SignalModel]],
Expand All @@ -99,12 +111,11 @@ def _send_signals_by_machine_id(
for machine_id in signals_by_machineid.keys()
]

retry_machines_to_process_attempts: List[MachineModel] = []
attempt_count = 0

while machines_to_process_attempts:
logging.info(f"attempt {attempt_count} to send signals")
retry_machines_to_process_attempts = []
retry_machines_to_process_attempts: List[MachineModel] = []
if attempt_count >= self.max_retries:
for machine_to_process in machines_to_process_attempts:
logging.error(
Expand Down Expand Up @@ -223,13 +234,15 @@ def _clear_all_signals(self):
self.storage.delete_signals(signals)

def _refresh_machine_token(self, machine: MachineModel) -> MachineModel:
scenarios = self.scenarios if self.force_config_scenario else machine.scenarios
machine.scenarios = (
self.scenarios if self.force_config_scenario else machine.scenarios
)
resp = self.http_client.post(
self._get_url(CAPI_WATCHER_LOGIN_ENDPOINT),
json={
"machine_id": machine.machine_id,
"password": machine.password,
"scenarios": scenarios.split(","),
"scenarios": machine.scenarios.split(","),
},
)
try:
Expand Down Expand Up @@ -279,7 +292,9 @@ def _ensure_machine_capi_registered(self, machine: MachineModel) -> MachineModel
return retrieved_machine

def _ensure_machine_capi_connected(self, machine: MachineModel) -> MachineModel:
if not has_valid_token(machine, self.latency_offset):
if not has_valid_token(
machine, self.latency_offset
) or not self.has_valid_scenarios(machine):
return self._refresh_machine_token(machine)
return machine

Expand Down

0 comments on commit 7955768

Please sign in to comment.