diff --git a/limacharlie/commands/_hive_shortcut.py b/limacharlie/commands/_hive_shortcut.py index 4659a99..55a226f 100644 --- a/limacharlie/commands/_hive_shortcut.py +++ b/limacharlie/commands/_hive_shortcut.py @@ -47,7 +47,12 @@ def make_hive_group(group_name: str, hive_name: str, noun_singular: str, noun_pl explain_list = f"List all {noun_plural} stored in the '{hive_name}' hive." explain_get = f"Get a specific {noun_singular} by its key name from the '{hive_name}' hive." - explain_set = f"Create or update {article} {noun_singular} in the '{hive_name}' hive. Provide data via --input-file (JSON/YAML) or stdin." + explain_set = ( + f"Create or update {article} {noun_singular} in the '{hive_name}' hive. " + f"Provide data via --input-file (JSON/YAML) or stdin. " + f"New hive records default to disabled — pass --enabled to create-and-enable in one shot, " + f"or include usr_mtd.enabled: true in the input file." + ) explain_delete = f"Delete {article} {noun_singular} from the '{hive_name}' hive. Requires --confirm for safety." @click.group(group_name) @@ -77,8 +82,12 @@ def get_cmd(ctx, key) -> None: @grp.command("set", help=f"Create or update {article} {noun_singular}.") @click.option("--key", required=True, help="Record key name.") @click.option("--input-file", type=click.Path(exists=True), default=None, help="JSON or YAML file with record data.") + @click.option( + "--enabled/--disabled", "enabled", default=None, + help=f"Set usr_mtd.enabled on the {noun_singular}. Overrides any value in the input file. Records default to disabled if neither this flag nor usr_mtd.enabled is provided.", + ) @pass_context - def set_cmd(ctx, key, input_file) -> None: + def set_cmd(ctx, key, input_file, enabled) -> None: if input_file: with open(input_file, "r") as f: content = f.read() @@ -108,6 +117,8 @@ def set_cmd(ctx, key, input_file) -> None: record = HiveRecord.from_raw(key, raw) else: record = HiveRecord(key, data=data) + if enabled is not None: + record.enabled = enabled org = _get_org(ctx) hive = Hive(org, hive_name) result = hive.set(record) diff --git a/limacharlie/commands/dr.py b/limacharlie/commands/dr.py index 906946e..a6bb265 100644 --- a/limacharlie/commands/dr.py +++ b/limacharlie/commands/dr.py @@ -254,9 +254,13 @@ def get(ctx, key, namespace) -> None: tags: [] comment: "rule description" +New D&R rules are created DISABLED by default for safety. Pass +--enabled to create-and-enable in one shot, or set usr_mtd.enabled +in the input. + Examples: - limacharlie dr set --key my-rule --input-file rule.yaml - cat rule.json | limacharlie dr set --key my-rule + limacharlie dr set --key my-rule --input-file rule.yaml --enabled + cat rule.json | limacharlie dr set --key my-rule --enabled limacharlie dr set --key my-rule --namespace managed --input-file rule.yaml IMPORTANT: Do not write D&R rules from scratch. Use @@ -278,8 +282,12 @@ def get(ctx, key, namespace) -> None: "--namespace", default=None, type=_NS_CHOICES, help="Namespace (default: general).", ) +@click.option( + "--enabled/--disabled", "enabled", default=None, + help="Set usr_mtd.enabled on the rule. Overrides any value in the input file. New rules default to disabled if neither this flag nor usr_mtd.enabled is provided.", +) @pass_context -def set_cmd(ctx, key, input_file, namespace) -> None: +def set_cmd(ctx, key, input_file, namespace, enabled) -> None: if input_file: with open(input_file, "r") as f: content = f.read() @@ -307,6 +315,9 @@ def set_cmd(ctx, key, input_file, namespace) -> None: else: record = HiveRecord(key, data=data) + if enabled is not None: + record.enabled = enabled + org = _get_org(ctx) hive = Hive(org, _hive_name(namespace)) result = hive.set(record) diff --git a/limacharlie/commands/hive.py b/limacharlie/commands/hive.py index 837e865..7a7cf1f 100644 --- a/limacharlie/commands/hive.py +++ b/limacharlie/commands/hive.py @@ -210,12 +210,15 @@ def get(ctx, hive_name, key) -> None: --input-file or from stdin if no file is specified. The input should be a JSON or YAML document. +New hive records are created DISABLED by default. Pass --enabled to +create-and-enable in one shot, or set usr_mtd.enabled in the input. + Full record format (YAML): data: key: value # payload varies by hive type usr_mtd: - enabled: true # optional, default true + enabled: true # optional, default false on new records expiry: 0 # optional, unix epoch (0 = never) tags: # optional - my-tag @@ -225,6 +228,9 @@ def get(ctx, hive_name, key) -> None: If the input has no "data" key, the entire input is treated as the record data payload. +The --enabled/--disabled flag, when given, overrides any value in +the input file's usr_mtd.enabled. + Data payload examples per hive type: secret: {secret: "my-api-key"} yara: {rule: "rule MyRule { ... }"} @@ -233,10 +239,10 @@ def get(ctx, hive_name, key) -> None: Examples: echo '{"data": {"key": "value"}}' | limacharlie hive set \\ - --hive-name lookup --key my-lookup + --hive-name lookup --key my-lookup --enabled limacharlie hive set --hive-name lookup --key my-lookup \\ - --input-file record.yaml + --input-file record.yaml --enabled """ register_explain("hive.set", _EXPLAIN_SET) @@ -245,8 +251,12 @@ def get(ctx, hive_name, key) -> None: @click.option("--hive-name", required=True, help="Hive name.") @click.option("--key", required=True, help="Record key.") @click.option("--input-file", default=None, type=click.Path(exists=True), help="Path to record data (JSON or YAML). Reads stdin if omitted.") +@click.option( + "--enabled/--disabled", "enabled", default=None, + help="Set usr_mtd.enabled on the record. Overrides any value in the input file. New records default to disabled if neither this flag nor usr_mtd.enabled is provided.", +) @pass_context -def set_record(ctx, hive_name, key, input_file) -> None: +def set_record(ctx, hive_name, key, input_file, enabled) -> None: data = _load_input(input_file) if data is None: click.echo( @@ -260,6 +270,8 @@ def set_record(ctx, hive_name, key, input_file) -> None: org = _get_org(ctx) hive = Hive(org, hive_name) record = _record_from_input(key, data) + if enabled is not None: + record.enabled = enabled result = hive.set(record) if not ctx.obj.quiet: click.echo(f"Record '{key}' set in hive '{hive_name}'.") diff --git a/tests/unit/test_cli_ai_skill_memory.py b/tests/unit/test_cli_ai_skill_memory.py index 9982b67..2f3bca6 100644 --- a/tests/unit/test_cli_ai_skill_memory.py +++ b/tests/unit/test_cli_ai_skill_memory.py @@ -62,6 +62,74 @@ def test_delete_requires_confirm(self): result = CliRunner().invoke(cli, ["ai-skill", "delete", "--key", "x"]) assert result.exit_code != 0 + @patch("limacharlie.commands._hive_shortcut.Client") + @patch("limacharlie.commands._hive_shortcut.Organization") + @patch("limacharlie.commands._hive_shortcut.Hive") + def test_set_enabled_flag_creates_enabled(self, mock_hive_cls, _org, _client): + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + result = CliRunner().invoke( + cli, ["ai-skill", "set", "--key", "triage", "--enabled"], + input=json.dumps({"content": "..."}), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is True + + @patch("limacharlie.commands._hive_shortcut.Client") + @patch("limacharlie.commands._hive_shortcut.Organization") + @patch("limacharlie.commands._hive_shortcut.Hive") + def test_set_no_flag_leaves_enabled_unset(self, mock_hive_cls, _org, _client): + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + result = CliRunner().invoke( + cli, ["ai-skill", "set", "--key", "triage"], + input=json.dumps({"content": "..."}), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + # Without the flag, enabled stays None so the API default applies. + assert record.enabled is None + + @patch("limacharlie.commands._hive_shortcut.Client") + @patch("limacharlie.commands._hive_shortcut.Organization") + @patch("limacharlie.commands._hive_shortcut.Hive") + def test_set_enabled_flag_overrides_input_file_value(self, mock_hive_cls, _org, _client): + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + payload = {"data": {"content": "..."}, "usr_mtd": {"enabled": False}} + result = CliRunner().invoke( + cli, ["ai-skill", "set", "--key", "triage", "--enabled"], + input=json.dumps(payload), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is True + + @patch("limacharlie.commands._hive_shortcut.Client") + @patch("limacharlie.commands._hive_shortcut.Organization") + @patch("limacharlie.commands._hive_shortcut.Hive") + def test_set_disabled_flag_works(self, mock_hive_cls, _org, _client): + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + # Input file says enabled=true, --disabled forces it off. + payload = {"data": {"content": "..."}, "usr_mtd": {"enabled": True}} + result = CliRunner().invoke( + cli, ["ai-skill", "set", "--key", "triage", "--disabled"], + input=json.dumps(payload), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is False + # --------------------------------------------------------------------------- # ai-memory (custom commands with partial-merge payloads) diff --git a/tests/unit/test_cli_commands.py b/tests/unit/test_cli_commands.py index cdb34b2..b4f38fc 100644 --- a/tests/unit/test_cli_commands.py +++ b/tests/unit/test_cli_commands.py @@ -509,6 +509,103 @@ def test_shortcut_disable(self, mock_hive_cls, mock_org_cls, mock_client_cls): assert record.tags == ["keep-me"] +class TestHiveSetEnabledFlag: + """The `--enabled/--disabled` flag on the create/update commands lets + AIs and operators create-and-enable a record in one shot, so they don't + forget the separate `enable` step and end up with silently-disabled + records. + """ + + @patch("limacharlie.commands.hive.Client") + @patch("limacharlie.commands.hive.Organization") + @patch("limacharlie.commands.hive.Hive") + def test_hive_set_with_enabled_flag(self, mock_hive_cls, _org, _client): + import json as _json + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + result = CliRunner().invoke( + cli, ["hive", "set", "--hive-name", "lookup", "--key", "k", "--enabled"], + input=_json.dumps({"data": {"v": 1}}), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is True + + @patch("limacharlie.commands.hive.Client") + @patch("limacharlie.commands.hive.Organization") + @patch("limacharlie.commands.hive.Hive") + def test_hive_set_flag_overrides_input_file(self, mock_hive_cls, _org, _client): + import json as _json + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + payload = {"data": {"v": 1}, "usr_mtd": {"enabled": False}} + result = CliRunner().invoke( + cli, ["hive", "set", "--hive-name", "lookup", "--key", "k", "--enabled"], + input=_json.dumps(payload), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is True + + @patch("limacharlie.commands.hive.Client") + @patch("limacharlie.commands.hive.Organization") + @patch("limacharlie.commands.hive.Hive") + def test_hive_set_no_flag_preserves_input_file_value(self, mock_hive_cls, _org, _client): + import json as _json + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + payload = {"data": {"v": 1}, "usr_mtd": {"enabled": True}} + result = CliRunner().invoke( + cli, ["hive", "set", "--hive-name", "lookup", "--key", "k"], + input=_json.dumps(payload), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is True + + @patch("limacharlie.commands.dr.Client") + @patch("limacharlie.commands.dr.Organization") + @patch("limacharlie.commands.dr.Hive") + def test_dr_set_with_enabled_flag(self, mock_hive_cls, _org, _client): + import json as _json + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + rule = {"detect": {"event": "NEW_PROCESS", "op": "is", "path": "event/FILE_PATH", "value": "x"}, "respond": [{"action": "report", "name": "x"}]} + result = CliRunner().invoke( + cli, ["dr", "set", "--key", "my-rule", "--enabled"], + input=_json.dumps(rule), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is True + + @patch("limacharlie.commands.dr.Client") + @patch("limacharlie.commands.dr.Organization") + @patch("limacharlie.commands.dr.Hive") + def test_dr_set_no_flag_leaves_enabled_unset(self, mock_hive_cls, _org, _client): + import json as _json + mock_hive = MagicMock() + mock_hive.set.return_value = {"etag": "e2"} + mock_hive_cls.return_value = mock_hive + + rule = {"detect": {"event": "NEW_PROCESS", "op": "is", "path": "event/FILE_PATH", "value": "x"}, "respond": [{"action": "report", "name": "x"}]} + result = CliRunner().invoke( + cli, ["dr", "set", "--key", "my-rule"], + input=_json.dumps(rule), + ) + assert result.exit_code == 0, result.output + record = mock_hive.set.call_args[0][0] + assert record.enabled is None + + class TestSchemaCommands: @patch("limacharlie.commands.schema.Client") @patch("limacharlie.commands.schema.Organization")