pyzgrab2 1.0.0__py3-none-any.whl

pyzgrab2/__init__.py

@@ -0,0 +1,27 @@
+from .client import ZGrab2
+from .exceptions import (
+    ZGrab2ConfigurationError,
+    ZGrab2Error,
+    ZGrab2NotFoundError,
+    ZGrab2OutputError,
+    ZGrab2ProcessError,
+    ZGrab2TimeoutError,
+)
+from .models import ScanResult, Target
+from .options import HTTPOptions, ScanOptions, SSHOptions, TLSOptions
+
+__all__ = [
+    'HTTPOptions',
+    'SSHOptions',
+    'ScanOptions',
+    'ScanResult',
+    'TLSOptions',
+    'Target',
+    'ZGrab2',
+    'ZGrab2ConfigurationError',
+    'ZGrab2Error',
+    'ZGrab2NotFoundError',
+    'ZGrab2OutputError',
+    'ZGrab2ProcessError',
+    'ZGrab2TimeoutError',
+]

pyzgrab2/_process.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import threading
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from pathlib import Path
+
+from .exceptions import ZGrab2ProcessError, ZGrab2TimeoutError
+
+
+def run_zgrab2(
+    executable: str | Path,
+    arguments: Sequence[str],
+    input_lines: Iterable[str],
+    *,
+    timeout: float | None = None,
+    cwd: str | Path | None = None,
+    env: Mapping[str, str] | None = None,
+) -> Iterator[str]:
+    '''Run zgrab2 and yield non-blank stdout lines as they arrive.'''
+    if timeout is not None and timeout <= 0:
+        raise ValueError('timeout must be greater than zero')
+
+    command = [str(executable), *arguments]
+    process_env = None if env is None else {**os.environ, **env}
+    try:
+        process = subprocess.Popen(
+            command,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            encoding='utf-8',
+            bufsize=1,
+            cwd=cwd,
+            env=process_env,
+        )
+    except OSError as error:
+        raise ZGrab2ProcessError(
+            f'Could not start zgrab2: {error}', command=command,
+        ) from error
+
+    if process.stdin is None or process.stdout is None or process.stderr is None:
+        process.kill()
+        process.wait()
+        raise ZGrab2ProcessError('Could not create pipes for zgrab2', command=command)
+
+    stdin = process.stdin
+    stdout = process.stdout
+    stderr_pipe = process.stderr
+    stderr_lines: list[str] = []
+    writer_errors: list[Exception] = []
+    timed_out = threading.Event()
+
+    def write_input() -> None:
+        try:
+            for line in input_lines:
+                stdin.write(line.rstrip('\r\n') + '\n')
+            stdin.flush()
+        except BrokenPipeError:
+            pass
+        except Exception as error:  # noqa: BLE001 - preserve producer failures
+            writer_errors.append(error)
+        finally:
+            stdin.close()
+
+    def read_stderr() -> None:
+        stderr_lines.extend(stderr_pipe)
+
+    def expire() -> None:
+        timed_out.set()
+        if process.poll() is None:
+            process.terminate()
+
+    writer_thread = threading.Thread(target=write_input, name='pyzgrab2-stdin', daemon=True)
+    stderr_thread = threading.Thread(target=read_stderr, name='pyzgrab2-stderr', daemon=True)
+    timer = threading.Timer(timeout, expire) if timeout is not None else None
+    writer_thread.start()
+    stderr_thread.start()
+    if timer is not None:
+        timer.start()
+
+    try:
+        for line in stdout:
+            if line.strip():
+                yield line
+
+        writer_thread.join()
+        stderr_thread.join()
+        return_code = process.wait()
+        stderr = ''.join(stderr_lines).strip()
+
+        if timed_out.is_set():
+            raise ZGrab2TimeoutError(
+                f'zgrab2 exceeded timeout of {timeout} seconds',
+                command=command,
+                returncode=return_code,
+                stderr=stderr,
+            )
+        if writer_errors:
+            raise ZGrab2ProcessError(
+                f'Failed while writing targets to zgrab2: {writer_errors[0]}',
+                command=command,
+                returncode=return_code,
+                stderr=stderr,
+            ) from writer_errors[0]
+        if return_code != 0:
+            message = f'zgrab2 exited with status {return_code}'
+            if stderr:
+                message += f': {stderr}'
+            raise ZGrab2ProcessError(
+                message, command=command, returncode=return_code, stderr=stderr,
+            )
+    finally:
+        if timer is not None:
+            timer.cancel()
+        if process.poll() is None:
+            process.terminate()
+            try:
+                process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait()
+        writer_thread.join(timeout=1)
+        stderr_thread.join(timeout=1)

pyzgrab2/client.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+import ipaddress
+import json
+import shutil
+import subprocess
+from collections.abc import Iterable, Iterator, Mapping
+from pathlib import Path
+
+from ._process import run_zgrab2
+from .exceptions import (
+    ZGrab2ConfigurationError,
+    ZGrab2NotFoundError,
+    ZGrab2OutputError,
+    ZGrab2ProcessError,
+)
+from .models import ScanResult, Target
+from .options import HTTPOptions, ScanOptions, SSHOptions, TLSOptions
+
+TargetInput = str | Target
+
+
+class ZGrab2:
+    '''Python interface to a local zgrab2 executable.'''
+
+    def __init__(self, executable: str | Path = 'zgrab2') -> None:
+        executable_path = shutil.which(str(executable))
+        if executable_path is None:
+            raise ZGrab2NotFoundError(f'Could not find the zgrab2 executable: {executable}')
+        self.executable = executable_path
+
+    def scan(
+        self,
+        targets: Iterable[TargetInput],
+        *,
+        module: str,
+        port: int | None = None,
+        module_args: Iterable[str] = (),
+        options: ScanOptions | None = None,
+        timeout: float | None = None,
+        cwd: str | Path | None = None,
+        env: Mapping[str, str] | None = None,
+    ) -> Iterator[ScanResult]:
+        '''Stream results from any zgrab2 module.'''
+        if not module or module.startswith('-'):
+            raise ZGrab2ConfigurationError('module must be a non-empty command name')
+        common = options or ScanOptions()
+        if port is not None and common.port is not None and port != common.port:
+            raise ZGrab2ConfigurationError('port conflicts with options.port')
+        if port is not None:
+            common = ScanOptions(
+                port=port,
+                senders=common.senders,
+                connections_per_host=common.connections_per_host,
+                connect_timeout=common.connect_timeout,
+                target_timeout=common.target_timeout,
+                read_limit_per_host=common.read_limit_per_host,
+                dns_resolvers=common.dns_resolvers,
+                dns_resolution_timeout=common.dns_resolution_timeout,
+                server_rate_limit=common.server_rate_limit,
+                resolve_ipv4=common.resolve_ipv4,
+                resolve_ipv6=common.resolve_ipv6,
+                debug=common.debug,
+                flush=common.flush,
+                verbose=common.verbose,
+                trigger=common.trigger,
+                extra_args=common.extra_args,
+            )
+        arguments = [module, *common.to_args(), *module_args]
+        input_lines = (self._normalize_target(target).to_csv_line() for target in targets)
+
+        for line_number, line in enumerate(
+            run_zgrab2(
+                executable=self.executable,
+                arguments=arguments,
+                input_lines=input_lines,
+                timeout=timeout,
+                cwd=cwd,
+                env=env,
+            ),
+            start=1,
+        ):
+            try:
+                raw = json.loads(line)
+            except (json.JSONDecodeError, UnicodeDecodeError) as error:
+                raise ZGrab2OutputError(
+                    f"Invalid JSON on zgrab2 output line {line_number}",
+                ) from error
+            if not isinstance(raw, dict):
+                raise ZGrab2OutputError(
+                    f"Expected a JSON object on zgrab2 output line {line_number}",
+                )
+            yield self._parse_result(raw, module, line_number=line_number)
+
+    def http(
+        self,
+        targets: Iterable[TargetInput],
+        *,
+        http: HTTPOptions | None = None,
+        options: ScanOptions | None = None,
+        timeout: float | None = None,
+    ) -> Iterator[ScanResult]:
+        '''Scan targets with zgrab2's HTTP module.'''
+        return self.scan(
+            targets,
+            module='http',
+            module_args=(http or HTTPOptions()).to_args(),
+            options=options,
+            timeout=timeout,
+        )
+
+    def tls(
+        self,
+        targets: Iterable[TargetInput],
+        *,
+        tls: TLSOptions | None = None,
+        options: ScanOptions | None = None,
+        timeout: float | None = None,
+    ) -> Iterator[ScanResult]:
+        '''Scan targets with zgrab2's TLS module.'''
+        return self.scan(
+            targets,
+            module='tls',
+            module_args=(tls or TLSOptions()).to_args(),
+            options=options,
+            timeout=timeout,
+        )
+
+    def ssh(
+        self,
+        targets: Iterable[TargetInput],
+        *,
+        ssh: SSHOptions | None = None,
+        options: ScanOptions | None = None,
+        timeout: float | None = None,
+    ) -> Iterator[ScanResult]:
+        '''Scan targets with zgrab2's SSH module.'''
+        return self.scan(
+            targets,
+            module='ssh',
+            module_args=(ssh or SSHOptions()).to_args(),
+            options=options,
+            timeout=timeout,
+        )
+
+
+    def module_help(self, module: str) -> str:
+        '''Return zgrab2's help text for a module.'''
+        if not module or module.startswith('-'):
+            raise ZGrab2ConfigurationError('module must be a non-empty command name')
+        return self._run_info_command(module, '--help')
+
+    def _run_info_command(self, *arguments: str) -> str:
+        command = [self.executable, *arguments]
+        try:
+            completed = subprocess.run(
+                command, capture_output=True, text=True, encoding='utf-8', check=False,
+            )
+        except OSError as error:
+            raise ZGrab2ProcessError(
+                f'Could not start zgrab2: {error}', command=command,
+            ) from error
+        if completed.returncode != 0:
+            stderr = completed.stderr.strip()
+            message = f"zgrab2 exited with status {completed.returncode}"
+            if stderr:
+                message += f": {stderr}"
+            raise ZGrab2ProcessError(
+                message,
+                command=command,
+                returncode=completed.returncode,
+                stderr=stderr,
+            )
+        return completed.stdout
+
+    @staticmethod
+    def _normalize_target(target: TargetInput) -> Target:
+        if isinstance(target, Target):
+            return target
+        try:
+            ipaddress.ip_network(target, strict=False)
+        except ValueError:
+            return Target(domain=target)
+        return Target(ip=target)
+
+    @staticmethod
+    def _parse_result(
+        raw: dict[str, object],
+        module: str,
+        *,
+        line_number: int | None = None,
+    ) -> ScanResult:
+        data = raw.get('data', {})
+        if not isinstance(data, dict):
+            raise ZGrab2OutputError(_unexpected_output('data must be an object', line_number))
+        module_output = data.get(module, {})
+        if module_output is None:
+            module_output = {}
+        if not isinstance(module_output, dict):
+            raise ZGrab2OutputError(
+                _unexpected_output(f"data.{module} must be an object", line_number),
+            )
+        return ScanResult(
+            ip=_optional_string(raw.get('ip'), 'ip', line_number),
+            domain=_optional_string(raw.get('domain'), 'domain', line_number),
+            tag=_optional_string(raw.get('tag'), 'tag', line_number),
+            module=module,
+            status=_optional_string(module_output.get('status'), 'status', line_number),
+            data=module_output,
+            raw=raw,
+        )
+
+
+def _optional_string(value: object, field: str, line_number: int | None) -> str | None:
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise ZGrab2OutputError(_unexpected_output(f'{field} must be a string', line_number))
+    return value
+
+
+def _unexpected_output(message: str, line_number: int | None) -> str:
+    suffix = f" on zgrab2 output line {line_number}" if line_number is not None else ''
+    return f"Unexpected zgrab2 output{suffix}: {message}"

pyzgrab2/exceptions.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+
+class ZGrab2Error(Exception):
+    '''Base exception for pyzgrab2.'''
+
+
+class ZGrab2NotFoundError(ZGrab2Error):
+    '''The zgrab2 executable could not be found.'''
+
+
+class ZGrab2ConfigurationError(ZGrab2Error, ValueError):
+    '''The requested scan configuration is invalid.'''
+
+
+class ZGrab2ProcessError(ZGrab2Error):
+    '''The zgrab2 process exited unsuccessfully.'''
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        command: Sequence[str] = (),
+        returncode: int | None = None,
+        stderr: str = '',
+    ) -> None:
+        super().__init__(message)
+        self.command = tuple(command)
+        self.returncode = returncode
+        self.stderr = stderr
+
+
+class ZGrab2TimeoutError(ZGrab2ProcessError, TimeoutError):
+    '''The zgrab2 process exceeded its allowed runtime.'''
+
+
+class ZGrab2OutputError(ZGrab2Error):
+    '''zgrab2 produced invalid or unexpected JSON output.'''

pyzgrab2/models.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import csv
+import io
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+from .exceptions import ZGrab2ConfigurationError
+
+
+def _validate_port(port: int | None) -> None:
+    if port is not None and not 1 <= port <= 65535:
+        raise ZGrab2ConfigurationError('Port must be between 1 and 65535')
+
+
+@dataclass(frozen=True, slots=True)
+class Target:
+    '''A zgrab2 input target in ``IP,DOMAIN,TAG,PORT`` form.'''
+
+    ip: str | None = None
+    domain: str | None = None
+    tag: str | None = None
+    port: int | None = None
+
+    def __post_init__(self) -> None:
+        if not self.ip and not self.domain:
+            raise ZGrab2ConfigurationError('Target requires an IP address, CIDR, or domain')
+
+        for value in (self.ip, self.domain, self.tag):
+            if value is not None and ('\n' in value or '\r' in value):
+                raise ZGrab2ConfigurationError('Target fields cannot contain newlines')
+
+        _validate_port(self.port)
+
+    def to_csv_line(self) -> str:
+        output = io.StringIO(newline='')
+        csv.writer(output, lineterminator='').writerow(
+            (self.ip or '', self.domain or '', self.tag or '', self.port or ''),
+        )
+        return output.getvalue()
+
+
+@dataclass(frozen=True, slots=True)
+class ScanResult:
+    '''One normalized zgrab2 JSON-lines result.'''
+
+    ip: str | None
+    domain: str | None
+    module: str
+    status: str | None
+    data: Mapping[str, Any] = field(default_factory=dict)
+    raw: Mapping[str, Any] = field(default_factory=dict)
+    tag: str | None = None
+
+    @property
+    def protocol(self) -> str | None:
+        value = self.data.get('protocol')
+        return value if isinstance(value, str) else None
+
+    @property
+    def port(self) -> int | None:
+        value = self.data.get('port')
+        return value if isinstance(value, int) and not isinstance(value, bool) else None
+
+    @property
+    def result(self) -> Mapping[str, Any] | None:
+        value = self.data.get('result')
+        return value if isinstance(value, Mapping) else None
+
+    @property
+    def error(self) -> str | None:
+        value = self.data.get('error')
+        return value if isinstance(value, str) else None
+
+    @property
+    def timestamp(self) -> datetime | None:
+        value = self.data.get('timestamp')
+        if not isinstance(value, str):
+            return None
+        try:
+            return datetime.fromisoformat(value.replace('Z', '+00:00'))
+        except ValueError:
+            return None