ptopenvas 0.0.1__py3-none-any.whl

ptopenvas/_version.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

ptopenvas/modules/gvm_setup.py

@@ -0,0 +1,237 @@
+import os
+import secrets
+import string
+import subprocess
+import sys
+import shutil
+import grp
+
+from gvm.transforms import EtreeCheckCommandTransform
+from gvm.connections import UnixSocketConnection
+from gvm.protocols.gmp import Gmp
+from gvm.errors import GvmError
+
+from ptlibs.app_dirs import AppDirs
+
+class GVMSetup:
+    
+    PASSWORD_FILE = os.path.join(AppDirs("ptopenvas").get_base_dir(), "gvm.pass")
+    
+    def __init__(self, args={}) -> None:    
+        self.args = args
+
+    def run(self) -> bool:
+        """
+        Main function to ensure dependencies and GVM setup.
+        Returns True if everything is ready.
+        """
+        # 1. Install GVM and ensure initialization
+        self.install_gvm() # ensure gvm is installed
+        self.ensure_gvmd_initialized() # run gvm-setup
+
+        # 2. Ensure GVMD service is running
+        self.run_gvmd_daemon()
+
+        # 3. Finalize permissions
+        self.finalize_permissions()
+        
+        # 4. Ensure penterep user exists (creates with random password if needed)
+        self.ensure_penterep_user()
+
+        return True
+    
+
+    def install_gvm(self) -> bool:
+        """
+        Installs missing GVM dependencies and runs gvm-setup.
+        Outputs all subprocess logs directly to the terminal.
+        """
+        required_packages = ["gvmd", "gvm-cli"] # Minimal required packages
+
+        missing = [pkg for pkg in required_packages if shutil.which(pkg) is None]
+
+        if missing:
+            print(f"Installing missing packages: {', '.join(missing)}")
+            try:
+                subprocess.run(["sudo", "apt", "update"], check=True)
+                subprocess.run(["sudo", "apt", "install", "-y"] + missing, check=True)
+                print("Installation complete.")
+            except subprocess.CalledProcessError as e:
+                print(f"Failed to install packages: {e}")
+                return False
+        else:
+            pass#print("All core GVM packages are already installed.")
+
+
+    def ensure_gvmd_initialized(self) -> bool:
+        """
+        Ensure GVMD is initialized by running `gvm-setup` if needed.
+
+        This checks for the GVMD Unix socket at `/run/gvmd/gvmd.sock` and
+        runs `gvm-setup` only when initialization appears required.
+        """
+        # 1) Check if GVMD database exists
+        try:
+            db_check = subprocess.run(
+                ["sudo", "-u", "postgres", "psql", "-tAc", "SELECT 1 FROM pg_database WHERE datname='gvmd'"],
+                #["-u", "postgres", "psql", "-tAc", "SELECT 1 FROM pg_database WHERE datname='gvmd'"],
+                capture_output=True, text=True, check=True
+            )
+            if db_check.stdout.strip() == "1":
+                #print("GVM setup was already completed previously.")
+                return True
+        except Exception:
+            pass
+
+        # If database missing -> setup never ran
+        try:
+            print("Running gvm-setup...")
+            subprocess.run(["sudo", "gvm-setup"], check=True)
+            print("GVMD initialization completed successfully.")
+            return True
+        except subprocess.CalledProcessError as e:
+            print(f"GVMD initialization failed: {e}")
+            return False
+
+    def run_gvmd_daemon(self):
+        """
+        Ensure that the GVMD service is running. 
+        Starts it if it is not active.
+        """
+        try:
+            #status = subprocess.run(["systemctl", "is-active", "--quiet", "gvmd"])
+            status = subprocess.run(["systemctl", "is-active", "--quiet", "gvmd"])
+            if status.returncode != 0:
+                print("GVMD service is not running. Starting with gvm-start...")
+                subprocess.run(["sudo", "gvm-start"], check=True)
+                print("GVM stack started.")
+            else:
+                pass
+                #print("GVMD service is already running.")
+        except subprocess.CalledProcessError as e:
+            print(f"Failed to start GVMD service: {e}")
+
+    def ensure_penterep_user(self):
+        """
+        Ensure gvmd service is running and the 'penterep' admin user exists.
+        If the user does not exist, create it with a random password stored in PASSWORD_FILE.
+        If the provided admin password is wrong, it will be reset automatically.
+        """
+
+        # 1. Define admin username and get password
+        USERNAME = "penterep"
+        PASSWORD = self.get_password() # "admin"
+
+        self._create_user_if_needed(username=USERNAME, password=PASSWORD)
+
+        # 2. Determine password for penterep
+        if not PASSWORD:
+            alphabet = string.ascii_letters + string.digits
+            PASSWORD = ''.join(secrets.choice(alphabet) for _ in range(24))
+            with open(self.PASSWORD_FILE, "w") as f:
+                f.write(PASSWORD)
+
+        # 3. Connect to gvmd and create user if missing
+        conn = UnixSocketConnection(path="/run/gvmd/gvmd.sock")
+        try:
+            with Gmp(conn, transform=EtreeCheckCommandTransform()) as gmp:
+                try:
+                    gmp.authenticate(USERNAME, PASSWORD)
+                except GvmError as e:
+                    try:
+                        subprocess.run([
+                            "sudo", "-u", "_gvm", "gvmd",
+                            "--user=penterep",
+                            f"--new-password={PASSWORD}"
+                        ], check=True, capture_output=True)
+                        #print(f"Reset penterep password to '{PASSWORD}'")
+                    except Exception as e:
+                        print(f"Failed to reset penterep password: {e}")
+                        return
+
+        except PermissionError:
+            print("Please reset your shell (log out and log back in) before running scripts that access GVMD.")
+        except GvmError as e:
+            print(f"Failed to create user: {e}")
+
+    def _create_user_if_needed(self, username, password):
+        r = subprocess.run([
+            "sudo", "-u", "_gvm", "gvmd",
+            "--get-users"
+            ], check=True, capture_output=True)
+
+        available_users = [u for u in r.stdout.decode().split("\n") if u]
+        #print("Available_users:\n   ", '\n    '.join(available_users))
+
+        if username not in available_users:
+            try:
+                # Create admin user
+                subprocess.run([
+                    "sudo", "-u", "_gvm", "gvmd",
+                    "--create-user="+username,
+                    "--password="+password,
+                ], check=True, capture_output=True)
+                print(f"Created admin user '{username}' with password '{password}'")
+            except Exception as e:
+                print(f"Failed to create user '{username}': {e}")
+                return
+
+
+    def get_password(self, generate: bool = False) -> str:
+        """
+        Return the stored password from PASSWORD_FILE.
+        If the file is missing or empty, or generate=True, create a new password, save it, and return it.
+        """
+        def generate_password() -> str:
+            charset = string.ascii_letters + string.digits
+            new_pass = ''.join(secrets.choice(charset) for _ in range(32))
+            with open(self.PASSWORD_FILE, "w") as f:
+                f.write(new_pass)
+            return new_pass
+
+        # Generate a new password if requested or if the file is missing/empty
+        if generate or not os.path.exists(self.PASSWORD_FILE):
+            return generate_password()
+
+        # Read existing password
+        with open(self.PASSWORD_FILE, "r") as f:
+            existing_pass = f.read().strip()
+
+        # If file is empty, generate a new password
+        if not existing_pass:
+            return generate_password()
+
+        return existing_pass
+
+
+    def finalize_permissions(self):
+        """
+        Add the current user to the _gvm group so the Python script can access the GVMD socket.
+        The user must reset their shell for the group change to take effect.
+        """
+        user = os.environ.get("USER")
+        if not user:
+            print("Could not determine the current user.")
+            return
+
+        # Check membership
+        try:
+            group = grp.getgrnam("_gvm")
+        except KeyError:
+            print("Group '_gvm' does not exist.")
+            return
+        
+        if user in group.gr_mem:
+            #print(f"User '{user}' is already a member of the _gvm group. No action taken.")
+            return
+
+        try:
+            subprocess.run(["sudo", "usermod", "-aG", "_gvm", user], check=True)
+            print(f"User '{user}' has been added to the _gvm group.")
+        except subprocess.CalledProcessError as e:
+            print(f"Failed to add user to _gvm group: {e}")
+
+
+if __name__ == "__main__":
+    gvm = GVMSetup(args={})
+    gvm.run()

ptopenvas/modules/helpers.py

@@ -0,0 +1,220 @@
+from lxml import etree
+import uuid
+from typing import List, Tuple, Union
+import time
+
+def _ensure_tree(maybe_xml: Union[str, etree._Element]) -> etree._Element:
+    """
+    Convert XML string to ElementTree if necessary.
+    """
+    if isinstance(maybe_xml, str):
+        return etree.fromstring(maybe_xml.encode())
+    return maybe_xml
+
+def _list_port_lists(gmp) -> List[Tuple[str, str, bool]]:
+    """
+    Return all port lists as a list of tuples: (id, name, predefined_flag)
+    """
+    root = _ensure_tree(gmp.get_port_lists())
+    items = []
+    for pl in root.xpath("//port_list"):
+        pid = pl.get("id") or pl.findtext("id")
+        name = pl.findtext("name") or ""
+        predefined_flag = pl.findtext("predefined") == "1"
+        items.append((pid, name, predefined_flag))
+    return items
+
+def _get_default_portlist_id(gmp, proto: str) -> str:
+    """
+    Return the ID of the default All TCP or All UDP port list.
+    Raises RuntimeError if not found.
+    """
+    proto = proto.lower()
+    defaults = {"tcp": "All IANA assigned TCP", "udp": "All UDP"}
+    root = _ensure_tree(gmp.get_port_lists())
+    node = root.xpath(f"//port_list[name='{defaults.get(proto, '')}']")
+    if not node:
+        raise RuntimeError(f"Default port list for {proto} not found")
+    return node[0].get("id")
+
+def _create_temp_portlist(gmp, ports: str) -> Tuple[str, str]:
+    """
+    Create a temporary port list in GVM for a scan.
+    ports: string like "22,80,443" or "1000-2000"
+    Returns: (portlist_id, name)
+    """
+    tmp_name = f"pt-{ports}-{time.strftime('%H:%M:%S')}"
+    created_xml = gmp.create_port_list(name=tmp_name, port_range=ports)
+    root = _ensure_tree(created_xml)
+    pid = root.findtext("id") or (root.xpath("//@id")[0] if root.xpath("//@id") else None)
+    if not pid:
+        raise RuntimeError("Failed to create temporary port list")
+    return pid, tmp_name
+
+def _cleanup_portlist(gmp, portlist_id: str) -> None:
+    """
+    Delete a temporary port list. Ignore errors.
+    """
+    try:
+        gmp.delete_port_list(portlist_id)
+    except Exception:
+        pass
+
+def get_default_scanner_id(gmp):
+    """
+    Retrieve the ID of the default OpenVAS scanner to assign tasks to.
+
+    Priority:
+      1. Scanner whose name contains "openvas" and status "Alive"
+      2. Any scanner whose name contains "openvas" (even if not alive)
+      3. Fallback to any alive scanner (if no OpenVAS exists)
+      4. Fallback to the first registered scanner
+
+    Args:
+        gmp: An authenticated GMP connection object.
+
+    Returns:
+        str: The ID of the selected scanner.
+
+    Raises:
+        RuntimeError: If no scanners are registered in GVMD.
+    """
+    root = _ensure_tree(gmp.get_scanners())
+    scanners = root.xpath("//scanner")
+    if not scanners:
+        raise RuntimeError("No scanners registered in gvmd")
+
+    # 1) Alive OpenVAS scanner
+    for s in scanners:
+        name = (s.findtext("name") or "").lower()
+        status = (s.findtext("status") or s.get("status") or "").lower()
+        if "openvas" in name and status == "alive":
+            return s.get("id") or s.findtext("id")
+
+    # 2) Any OpenVAS scanner (not necessarily alive)
+    for s in scanners:
+        name = (s.findtext("name") or "").lower()
+        if "openvas" in name:
+            return s.get("id") or s.findtext("id")
+
+    # 3) Any alive scanner
+    for s in scanners:
+        status = (s.findtext("status") or s.get("status") or "").lower()
+        if status == "alive":
+            return s.get("id") or s.findtext("id")
+
+    # 4) Fallback: first scanner
+    s = scanners[0]
+    return s.get("id") or s.findtext("id")
+
+
+def wait_for_report(gmp, task_id, timeout=300, interval=5, verbose=False):
+    """
+    Wait until a report is available for the given task.
+
+    Strategy:
+     - Poll gmp.get_task(task_id) and look for report id in common places:
+         * .//last_report/id
+         * .//report/id
+         * .//report (check @id)
+     - If task status becomes a finished state (Done/Stopped/Aborted/Failed) but no report id found,
+       query gmp.get_reports() and filter for reports whose task/id == task_id and return the newest.
+     - After timeout, do a final search in reports and raise RuntimeError if nothing found.
+
+    Args:
+        gmp: Authenticated GMP connection object.
+        task_id: ID of the task.
+        timeout: Seconds to wait before giving up.
+        interval: Poll interval in seconds.
+        verbose: If True, prints progress messages.
+
+    Returns:
+        report_id (str)
+
+    Raises:
+        RuntimeError: If no report is found within timeout.
+    """
+    waited = 0
+    finished_states = {"done", "stopped", "aborted", "cancelled", "failed"}
+    if verbose:
+        print(f"[wait_for_report] waiting for report for task {task_id} (timeout={timeout}s)")
+
+    while waited < timeout:
+        task_xml = gmp.get_task(task_id)
+        # check several common locations for report id
+        report_id = (
+            task_xml.findtext(".//last_report/id")
+            or task_xml.findtext(".//report/id")
+        )
+
+        # sometimes report is present as a <report id="..."> element
+        if not report_id:
+            report_nodes = task_xml.xpath(".//report")
+            if report_nodes:
+                # try attribute 'id' first, fallback to child <id>
+                first = report_nodes[0]
+                report_id = first.get("id") or first.findtext("id")
+
+        status = (task_xml.findtext(".//status") or "").strip().lower()
+        if verbose:
+            print(f"[wait_for_report] waited={waited}s status='{status}' report_id='{report_id}'")
+
+        if report_id:
+            if verbose:
+                print(f"[wait_for_report] found report id on task: {report_id}")
+            return report_id
+
+        # If the task is finished but no report id on task, search reports list
+        if status in finished_states:
+            if verbose:
+                print(f"[wait_for_report] task in finished state '{status}' but no report on task; searching reports table...")
+            # search reports for this task
+            candidate = _find_latest_report_for_task(gmp, task_id, verbose=verbose)
+            if candidate:
+                return candidate
+            # else continue waiting a little in case the report is still being generated/attached
+
+        time.sleep(interval)
+        waited += interval
+
+    # Final attempt after timeout
+    if verbose:
+        print(f"[wait_for_report] timeout reached ({timeout}s). final search in reports...")
+    candidate = _find_latest_report_for_task(gmp, task_id, verbose=verbose)
+    if candidate:
+        return candidate
+
+    raise RuntimeError(f"No report found for task {task_id} after waiting {timeout} seconds")
+
+
+def _find_latest_report_for_task(gmp, task_id, verbose=False):
+    """
+    Search all reports and return the newest report id for the given task_id, or None.
+    Uses gmp.get_reports() and XPath; returns the first match or the newest by creation_time if available.
+    """
+    reports_xml = gmp.get_reports()
+    reports = reports_xml.xpath(f"//report[task/id='{task_id}']")
+    if not reports:
+        if verbose:
+            print(f"[_find_latest_report_for_task] no reports found for task {task_id}")
+        return None
+
+    # prefer attribute id, else child <id>; prefer newest by creation_time if present
+    def extract_info(rnode):
+        rid = rnode.get("id") or rnode.findtext("id")
+        # try to get creation time if available (many report formats include it)
+        ctime = rnode.findtext("creation_time") or rnode.findtext("report/creation_time") or ""
+        return (rid, ctime)
+
+    infos = [extract_info(r) for r in reports]
+    # if any creation times present, sort by them (lexicographic ISO-8601 is fine)
+    if any(info[1] for info in infos):
+        infos = sorted(infos, key=lambda x: x[1] or "", reverse=True)
+    else:
+        # fallback: return first report node's id
+        infos = infos
+
+    chosen = infos[0][0] if infos else None
+    if verbose:
+        print(f"[_find_latest_report_for_task] candidate reports: {[i[0] for i in infos]}; chosen={chosen}")
+    return chosen