Skip to content

Commit

Permalink
fix: use proper xpaths, convert types
Browse files Browse the repository at this point in the history
  • Loading branch information
linuxdaemon committed May 12, 2024
1 parent 48781b0 commit 62455b8
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 11 deletions.
32 changes: 21 additions & 11 deletions src/gvm_sync_targets/cli/__init__.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
# SPDX-FileCopyrightText: 2024-present linuxdaemon <linuxdaemon.irc@gmail.com>
#
# SPDX-License-Identifier: MIT
from typing import TYPE_CHECKING, Optional, TextIO
from typing import TYPE_CHECKING, TextIO, cast

import click
from gvm.connections import UnixSocketConnection
from gvm.protocols.gmp import Gmp
from gvm.transforms import EtreeCheckCommandTransform

from gvm_sync_targets import __version__
from gvm_sync_targets.util import to_str

if TYPE_CHECKING:
from lxml.etree import ElementBase
Expand Down Expand Up @@ -37,27 +38,36 @@ def gvm_sync_targets(username: str, password: str, hosts_file: TextIO) -> None:
if target is None:
# Create target
new_target = gmp.create_target("All Hosts")
elif target.xpath("in_use/text()") == ["1"]:
new_target = gmp.clone_target(target.attrib["id"])
elif target.xpath("sum(in_use/text()) > 0"):
new_target = gmp.clone_target(to_str(target.attrib["id"]))
else:
new_target = target

new_target_id: str = to_str(new_target.attrib["id"])
# Set hosts
# gmp.modify_target(new_target.attrib['id'], hosts=hosts)

# Replace old target in tasks
if target is not None and target is not new_target:
tasks = gmp.get_tasks().findall("task")
if tasks is None:
tasks = []
old_target_id = to_str(target.attrib["id"])
tasks = cast("ElementBase", gmp.get_tasks()).findall(
f"task[target='{old_target_id}']"
)

for task in tasks:
print(task)
print(task.attrib)
for child in task:
print(child)
print(child.attrib)
gmp.modify_task(
to_str(task.attrib["id"]), target_id=new_target_id
)

# Ensure old target is unused
old_target: "ElementBase" = gmp.get_target(old_target_id)
old_name = old_target.xpath("name/text()[1]")
if old_target.xpath("sum(in_use/text()) > 0"):
msg = "target is still in use"
raise ValueError(msg)

# Delete old target
gmp.delete_target(old_target_id, ultimate=True)

# Rename new target
gmp.modify_target(new_target_id, name=old_name)
32 changes: 32 additions & 0 deletions src/gvm_sync_targets/util.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# SPDX-FileCopyrightText: 2024-present linuxdaemon <linuxdaemon.irc@gmail.com>
#
# SPDX-License-Identifier: MIT
from typing import Union


def to_str(text: Union[str, bytes, bytearray, memoryview]) -> str:
"""Convert bytes, bytearray, memoryview, or str to str.
Args:
text: Input value
Returns:
decoded output
Examples:
>>> to_str("aaa")
'aaa'
>>> to_str(b'bb')
'bb'
>>> to_str(memoryview(b'bb'))
'bb'
>>> to_str(bytearray(b'cc'))
'cc'
"""
if isinstance(text, (bytes, bytearray)):
return text.decode()

if isinstance(text, memoryview):
return to_str(text.tobytes())

return text

0 comments on commit 62455b8

Please sign in to comment.