Skip to content

add: support scheduler max delay in agent GMP commands #1395

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.

Sign up for GitHub

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

Open
ozgen wants to merge 1 commit into
base: main
Choose a base branch
from
Open

add: support scheduler max delay in agent GMP commands #1395

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 6 additions & 0 deletions gvm/protocols/gmp/_gmpnext.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,7 @@ def create_agent_group(
agent_ids: list[str],
scheduler_cron_time: str,
*,
max_delay_seconds: int | None = None,
comment: str | None = None,
) -> T:
"""Create a new agent group.
Expand All @@ -226,6 +227,7 @@ def create_agent_group(
name: Name of the new agent group.
agent_ids: List of agent UUIDs to include in the group (required).
scheduler_cron_time: Scheduler cron to use.
max_delay_seconds: Max delay seconds to use.
comment: Optional comment for the group.

Raises:
Expand All @@ -237,6 +239,7 @@ def create_agent_group(
comment=comment,
agent_ids=agent_ids,
scheduler_cron_time=scheduler_cron_time,
max_delay_seconds=max_delay_seconds,
)
)

Expand All @@ -248,6 +251,7 @@ def modify_agent_group(
name: str | None = None,
comment: str | None = None,
agent_ids: list[str] | None = None,
max_delay_seconds: int | None = None,
) -> T:
"""Modify an existing agent group.

Expand All @@ -257,6 +261,7 @@ def modify_agent_group(
scheduler_cron_time: Scheduler cron to use.
comment: Optional comment for the group.
agent_ids: Optional list of agent UUIDs to set for the group.
max_delay_seconds: Optional maximum delay for agent groups.

Raises:
RequiredArgument: If agent_group_id is not provided.
Expand All @@ -268,6 +273,7 @@ def modify_agent_group(
comment=comment,
agent_ids=agent_ids,
scheduler_cron_time=scheduler_cron_time,
max_delay_seconds=max_delay_seconds,
)
)

Expand Down
9 changes: 9 additions & 0 deletions gvm/protocols/gmp/requests/next/_agent_groups.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ def create_agent_group(
agent_ids: list[str],
scheduler_cron_time: str,
*,
max_delay_seconds: int | None = None,
comment: str | None = None,
) -> Request:
"""Create a new agent group.
Expand All @@ -26,6 +27,7 @@ def create_agent_group(
name: Name of the new agent group.
agent_ids: List of agent UUIDs to include in the group (required).
scheduler_cron_time: Cron-like time to schedule new agent groups (required).
max_delay_seconds: Max delay seconds to schedule the new agent group (optional).
comment: Optional comment for the group.

Raises:
Expand All @@ -50,6 +52,9 @@ def create_agent_group(
cmd.add_element("name", name)
cmd.add_element("scheduler_cron_time", scheduler_cron_time)

if max_delay_seconds is not None:
cmd.add_element("max_delay_seconds", str(max_delay_seconds))

if comment:
cmd.add_element("comment", comment)

Expand Down Expand Up @@ -89,6 +94,7 @@ def modify_agent_group(
name: str | None = None,
comment: str | None = None,
agent_ids: list[str] | None = None,
max_delay_seconds: int | None = None,
) -> Request:
"""Modify an existing agent group.

Expand All @@ -98,6 +104,7 @@ def modify_agent_group(
name: Optional new name for the group.
comment: Optional comment for the group.
agent_ids: Optional list of agent UUIDs to set for the group.
max_delay_seconds: Max delay seconds to schedule the new agent group.

Raises:
RequiredArgument: If agent_group_id is not provided.
Expand Down Expand Up @@ -127,6 +134,8 @@ def modify_agent_group(
agents_element = cmd.add_element("agents")
for agent_id in agent_ids:
agents_element.add_element("agent", attrs={"id": agent_id})
if max_delay_seconds is not None:
cmd.add_element("max_delay_seconds", str(max_delay_seconds))

return cmd

Expand Down
79 changes: 61 additions & 18 deletions gvm/protocols/gmp/requests/next/_agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,20 +73,40 @@ def valid_value(d: Mapping[str, Any], key: str, path: str) -> Any:
valid_value(se, "indexer_dir_depth", "agent_script_executor.")

sched = se.get("scheduler_cron_time")
if isinstance(sched, str):
items = [sched]
elif isinstance(sched, Sequence) and not isinstance(
sched, (str, bytes)

if (
not isinstance(sched, Sequence)
or isinstance(sched, (str, bytes))
or not sched
):
items = [str(x) for x in sched]
else:
items = []
if not items or any(not str(x).strip() for x in items):
raise RequiredArgument(
function=caller,
argument="config.agent_script_executor.scheduler_cron_time",
)

for item in sched:
if not isinstance(item, Mapping):
raise InvalidArgumentType(
function=caller,
argument="config.agent_script_executor.scheduler_cron_time",
arg_type="Mapping",
)

cron = item.get("cron")
max_delay_seconds = item.get("max_delay_seconds")

if not isinstance(cron, str) or not cron.strip():
raise RequiredArgument(
function=caller,
argument="config.agent_script_executor.scheduler_cron_time.cron",
)

if type(max_delay_seconds) is not int or max_delay_seconds < 0:
raise InvalidArgumentType(
function=caller,
argument="config.agent_script_executor.scheduler_cron_time.max_delay_seconds",
arg_type="non-negative integer",
)
# heartbeat
hb = valid_map(config, "heartbeat", "")
valid_value(hb, "interval_in_seconds", "heartbeat.")
Expand Down Expand Up @@ -162,7 +182,12 @@ def _append_agent_config(
"bulk_size": 2,
"bulk_throttle_time_in_ms": 300,
"indexer_dir_depth": 100,
"scheduler_cron_time": ["0 */12 * * *"],
"scheduler_cron_time": [
{
"cron": "0 */12 * * *",
"max_delay_seconds": 30,
},
],
},
"heartbeat": {
"interval_in_seconds": 300,
Expand Down Expand Up @@ -212,18 +237,21 @@ def _append_agent_config(
xml_se, "indexer_dir_depth", se.get("indexer_dir_depth")
)
sched = se.get("scheduler_cron_time")
if isinstance(sched, str):
sched_items = [sched]
else:
sched_items = list(sched or [])

if sched_items:
if sched:
xml_sched = xml_se.add_element(
"scheduler_cron_time",
attrs={"is_list": "1"},
)
for item in sched_items:
xml_sched.add_element("item", str(item))

for item in sched:
xml_item = xml_sched.add_element("item")
cls._add_element(xml_item, "cron", item["cron"])
cls._add_element(
xml_item,
"max_delay_seconds",
item["max_delay_seconds"],
)

# heartbeat
hb = config["heartbeat"]
Expand Down Expand Up @@ -311,7 +339,12 @@ def modify_agents(
"bulk_size": 2,
"bulk_throttle_time_in_ms": 300,
"indexer_dir_depth": 100,
"scheduler_cron_time": ["0 */12 * * *"], # str or list[str]
"scheduler_cron_time": [
{
"cron": "0 */12 * * *",
"max_delay_seconds": 30,
},
],
},
"heartbeat": {
"interval_in_seconds": 300,
Expand Down Expand Up @@ -393,7 +426,16 @@ def modify_agent_control_scan_config(
"bulk_size": 2,
"bulk_throttle_time_in_ms": 300,
"indexer_dir_depth": 100,
"scheduler_cron_time": ["0 */12 * * *"],
"scheduler_cron_time": [
{
"cron": "0 */12 * * *",
"max_delay_seconds": 30,
},
{
"cron": "0 23 * * *",
"max_delay_seconds": 60,
},
],
},
"heartbeat": {
"interval_in_seconds": 300,
Expand All @@ -410,6 +452,7 @@ def modify_agent_control_scan_config(
function=cls.modify_agent_control_scan_config.__name__,
argument="agent_control_id",
)

if not config_defaults:
raise RequiredArgument(
function=cls.modify_agent_control_scan_config.__name__,
Expand Down
34 changes: 34 additions & 0 deletions tests/protocols/gmpnext/entities/agent_groups/test_create_agent_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,14 @@ def test_create_agent_group(self):
agent_ids=["agent-1", "agent-2"],
comment="Sample comment",
scheduler_cron_time="0 */5 * * *",
max_delay_seconds=30,
)

self.connection.send.has_been_called_with(
b"<create_agent_group>"
b"<name>ExampleGroup</name>"
b"<scheduler_cron_time>0 */5 * * *</scheduler_cron_time>"
b"<max_delay_seconds>30</max_delay_seconds>"
b"<comment>Sample comment</comment>"
b'<agents><agent id="agent-1"/><agent id="agent-2"/></agents>'
b"</create_agent_group>"
Expand Down Expand Up @@ -57,3 +59,35 @@ def test_create_agent_group_without_comment(self):
b'<agents><agent id="a1"/><agent id="a2"/></agents>'
b"</create_agent_group>"
)

def test_create_agent_group_with_zero_max_delay(self):
self.gmp.create_agent_group(
name="GroupX",
agent_ids=["a1"],
scheduler_cron_time="0 */5 * * *",
max_delay_seconds=0,
)

self.connection.send.has_been_called_with(
b"<create_agent_group>"
b"<name>GroupX</name>"
b"<scheduler_cron_time>0 */5 * * *</scheduler_cron_time>"
b"<max_delay_seconds>0</max_delay_seconds>"
b'<agents><agent id="a1"/></agents>'
b"</create_agent_group>"
)

def test_create_agent_group_without_max_delay(self):
self.gmp.create_agent_group(
name="GroupX",
agent_ids=["a1"],
scheduler_cron_time="0 */5 * * *",
)

self.connection.send.has_been_called_with(
b"<create_agent_group>"
b"<name>GroupX</name>"
b"<scheduler_cron_time>0 */5 * * *</scheduler_cron_time>"
b'<agents><agent id="a1"/></agents>'
b"</create_agent_group>"
)
28 changes: 28 additions & 0 deletions tests/protocols/gmpnext/entities/agent_groups/test_modify_agent_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ def test_modify_agent_group(self):
name="NewName",
comment="Updated comment",
scheduler_cron_time="0 */5 * * *",
max_delay_seconds=30,
agent_ids=["agent-1", "agent-2"],
)

Expand All @@ -21,6 +22,7 @@ def test_modify_agent_group(self):
b"<name>NewName</name>"
b"<comment>Updated comment</comment>"
b'<agents><agent id="agent-1"/><agent id="agent-2"/></agents>'
b"<max_delay_seconds>30</max_delay_seconds>"
b"</modify_agent_group>"
)

Expand Down Expand Up @@ -83,3 +85,29 @@ def test_modify_agent_group_without_agent_ids(self):
b"<comment>Updated comment</comment>"
b"</modify_agent_group>"
)

def test_modify_agent_group_with_zero_max_delay(self):
self.gmp.modify_agent_group(
agent_group_id="group-1",
scheduler_cron_time="0 */5 * * *",
max_delay_seconds=0,
)

self.connection.send.has_been_called_with(
b'<modify_agent_group agent_group_id="group-1">'
b"<scheduler_cron_time>0 */5 * * *</scheduler_cron_time>"
b"<max_delay_seconds>0</max_delay_seconds>"
b"</modify_agent_group>"
)

def test_modify_agent_group_without_max_delay(self):
self.gmp.modify_agent_group(
agent_group_id="group-1",
scheduler_cron_time="0 */5 * * *",
)

self.connection.send.has_been_called_with(
b'<modify_agent_group agent_group_id="group-1">'
b"<scheduler_cron_time>0 */5 * * *</scheduler_cron_time>"
b"</modify_agent_group>"
)
Loading
Loading