From a45d4f1fce2221bd438b7c03ec0b3f560a1ff554 Mon Sep 17 00:00:00 2001 From: Tyson Smith Date: Thu, 18 Apr 2024 09:55:39 -0700 Subject: [PATCH] Add type hints to adapter --- grizzly/adapter/adapter.py | 113 +++++++++++--------- grizzly/adapter/no_op_adapter/__init__.py | 18 ++-- grizzly/adapter/no_op_adapter/test_no_op.py | 6 +- grizzly/main.py | 2 +- 4 files changed, 75 insertions(+), 64 deletions(-) diff --git a/grizzly/adapter/adapter.py b/grizzly/adapter/adapter.py index 924754c4..67cc54a6 100644 --- a/grizzly/adapter/adapter.py +++ b/grizzly/adapter/adapter.py @@ -3,8 +3,12 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. from abc import ABCMeta, abstractmethod from pathlib import Path +from typing import Any, Dict, Generator, Optional, Tuple +from grizzly.common.storage import TestCase from grizzly.common.utils import DEFAULT_TIME_LIMIT, HARNESS_FILE +from grizzly.target.target_monitor import TargetMonitor +from sapphire import ServerMap __all__ = ("Adapter", "AdapterError") __author__ = "Tyson Smith" @@ -19,21 +23,21 @@ class Adapter(metaclass=ABCMeta): """An Adapter is the interface between Grizzly and a fuzzer. A subclass must be created in order to add support for additional fuzzers. The Adapter is responsible for handling input/output data and executing the fuzzer. - It is expected that any processes launched or file created on file system + It is expected that any processes launched or files created on file system by the adapter will also be cleaned up by the adapter. - NOTE: Some methods must not be overloaded doing so will prevent Grizzly from + + NOTE: Some methods must not be overridden doing so will prevent Grizzly from operating correctly. Attributes: - _harness (Path): Harness file that will be used. - fuzz (dict): Available as a safe scratch pad for the end-user. - monitor (TargetMonitor): Used to provide Target status information to - the adapter. - name (str): Name of the adapter. - remaining (int): Can be used to indicate the number of TestCases - remaining to process. + _harness: Harness file that will be used. + fuzz: Available as a safe scratch pad for the end-user. + monitor: Used to provide Target status information to the adapter. + name: Name of the adapter. + remaining: Can be used to indicate the number of TestCases remaining to process. """ + IGNORE_FILES = ("desktop.ini", "thumbs.db") # Maximum iterations between Target relaunches (<1 use default) RELAUNCH = 0 # Maximum execution time per test (used as minimum timeout). The iteration is @@ -43,27 +47,27 @@ class Adapter(metaclass=ABCMeta): __slots__ = ("_harness", "fuzz", "monitor", "name", "remaining") - def __init__(self, name): + def __init__(self, name: str) -> None: assert isinstance(name, str) if not name: raise AdapterError("name must not be empty") if len(name.split()) != 1 or name.strip() != name: raise AdapterError("name must not contain whitespace") - self._harness = None - self.fuzz = {} - self.monitor = None + self._harness: Optional[bytes] = None + self.fuzz: Dict[str, Any] = {} + self.monitor: Optional[TargetMonitor] = None self.name = name - self.remaining = None + self.remaining: Optional[int] = None - def __enter__(self): + def __enter__(self) -> "Adapter": return self - def __exit__(self, *exc): + def __exit__(self, *exc: Any) -> None: self.cleanup() - def cleanup(self): + def cleanup(self) -> None: """Automatically called once at shutdown. Used internally by Grizzly. - *** DO NOT OVERLOAD! *** + *** DO NOT OVERRIDE! *** Args: None @@ -73,96 +77,99 @@ def cleanup(self): """ self.shutdown() - def enable_harness(self, file_path=HARNESS_FILE): + def enable_harness(self, path: Path = HARNESS_FILE) -> None: """Enable use of a harness during fuzzing. By default no harness is used. - *** DO NOT OVERLOAD! *** + *** DO NOT OVERRIDE! *** Args: - file_path (StrPath): HTML file to use as a harness. + path: HTML file to use as a harness. Returns: None """ - path = Path(file_path) - assert path.is_file(), f"missing harness file '{path.resolve()}'" self._harness = path.read_bytes() assert self._harness, f"empty harness file '{path.resolve()}'" - def get_harness(self): + def get_harness(self) -> Optional[bytes]: """Get the harness. Used internally by Grizzly. - *** DO NOT OVERLOAD! *** + *** DO NOT OVERRIDE! *** Args: None Returns: - bytes: The active harness. + The active harness data. """ return self._harness + # TODO: change return type from str to Path and accept Path instead of str @staticmethod - def scan_path(path, ignore=("desktop.ini", "thumbs.db"), recursive=False): + def scan_path( + path: str, + ignore: Tuple[str, ...] = IGNORE_FILES, + recursive: bool = False, + ) -> Generator[str, None, None]: """Scan a path and yield the files within it. This is available as a helper method. Args: - path (str): Path to file or directory. - ignore (iterable(str)): Files to ignore. - recursive (bool): Scan recursively into directories. + path: Path to file or directory. + ignore: File names to ignore. + recursive: Scan recursively into directories. Yields: - str: Absolute path to files. + Absolute path to files. """ - path = Path(path).resolve() - if path.is_dir(): - path_iter = path.rglob("*") if recursive else path.glob("*") + src = Path(path) + if src.is_dir(): + path_iter = src.rglob("*") if recursive else src.glob("*") for entry in path_iter: if not entry.is_file(): continue - if entry.name in ignore or entry.name.startswith("."): + if entry.name.lower() in ignore or entry.name.startswith("."): # skip ignored and hidden system files continue - yield str(entry) - elif path.is_file(): - yield str(path) + yield str(entry.resolve()) + elif src.is_file(): + yield str(src.resolve()) @abstractmethod - def generate(self, testcase, server_map): + def generate(self, testcase: TestCase, server_map: ServerMap) -> None: """Automatically called. Populate testcase here. Args: - testcase (TestCase): TestCase intended to be populated. - server_map (ServerMap): A ServerMap. + testcase: TestCase intended to be populated. + server_map: A ServerMap. Returns: None """ - def on_served(self, testcase, served): + def on_served(self, testcase: TestCase, served: Tuple[str, ...]) -> None: """Optional. Automatically called after a test case is successfully served. Args: - testcase (TestCase): TestCase that was served. - served (tuple(str)): Files served from testcase. + testcase: TestCase that was served. + served: Files served from testcase. Returns: None """ - def on_timeout(self, testcase, served): + def on_timeout(self, testcase: TestCase, served: Tuple[str, ...]) -> None: """Optional. Automatically called if timeout occurs while attempting to serve a test case. By default it calls `self.on_served()`. Args: - testcase (TestCase): TestCase that was served. - served (tuple(str)): Files served from testcase. + testcase: TestCase that was served. + served: Files served from testcase. Returns: None """ self.on_served(testcase, served) - def pre_launch(self): + def pre_launch(self) -> None: """Optional. Automatically called before launching the Target. Args: @@ -172,19 +179,19 @@ def pre_launch(self): None """ - def setup(self, input_path, server_map): + # TODO: update input_path type (str -> Path) + def setup(self, input_path: Optional[str], server_map: ServerMap) -> None: """Optional. Automatically called once at startup. Args: - input_path (str): Points to a file or directory passed by the user. - None is passed by default. - server_map (ServerMap): A ServerMap + input_path: File or directory passed by the user. + server_map: A ServerMap Returns: None """ - def shutdown(self): + def shutdown(self) -> None: """Optional. Automatically called once at shutdown. Args: diff --git a/grizzly/adapter/no_op_adapter/__init__.py b/grizzly/adapter/no_op_adapter/__init__.py index f056982f..28885871 100644 --- a/grizzly/adapter/no_op_adapter/__init__.py +++ b/grizzly/adapter/no_op_adapter/__init__.py @@ -2,7 +2,11 @@ # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. +from typing import Optional + from grizzly.adapter import Adapter +from grizzly.common.storage import TestCase +from sapphire import ServerMap __author__ = "Tyson Smith" __credits__ = ["Tyson Smith"] @@ -15,15 +19,15 @@ class NoOpAdapter(Adapter): NAME = "no-op" - def setup(self, _input, _server_map): + def setup(self, input_path: Optional[str], server_map: ServerMap) -> None: """Generate a static test case that calls `window.close()` when run. Normally this is done in generate() but since the test is static only do it once. Use the default harness to allow running multiple test cases in a row without closing the browser after each one. Args: - _input (str): Unused. - _server_map (sapphire.server_map.ServerMap): Unused. + _input: Unused. + _server_map: Unused. Returns: None @@ -38,7 +42,7 @@ def setup(self, _input, _server_map): b"" ) - def generate(self, testcase, _server_map): + def generate(self, testcase: TestCase, server_map: ServerMap) -> None: """The test case contents have been created now add the data to the TestCase. All TestCases require an entry point and the one expected by Grizzly @@ -46,10 +50,10 @@ def generate(self, testcase, _server_map): the test. Args: - testcase (grizzly.common.storage.TestCase): TestCase to be populated. - _server_map (sapphire.server_map.ServerMap): Unused in this example. + testcase: TestCase to be populated. + _server_map: Unused in this example. Returns: None """ - testcase.add_from_bytes(self.fuzz["test"], testcase.entry_point) + testcase.add_from_bytes(self.fuzz["test"], testcase.entry_point, required=True) diff --git a/grizzly/adapter/no_op_adapter/test_no_op.py b/grizzly/adapter/no_op_adapter/test_no_op.py index 954b8a56..31708750 100644 --- a/grizzly/adapter/no_op_adapter/test_no_op.py +++ b/grizzly/adapter/no_op_adapter/test_no_op.py @@ -11,8 +11,8 @@ def test_no_op_01(): """test a simple Adapter""" with NoOpAdapter("no-op") as adapter: adapter.setup(None, None) - with TestCase("a", adapter.name) as test: + with TestCase("a.html", adapter.name) as test: assert not test.data_size - assert "a" not in test + assert "a.html" not in test adapter.generate(test, None) - assert "a" in test + assert "a.html" in test diff --git a/grizzly/main.py b/grizzly/main.py index f4c21b1d..f902fe9d 100644 --- a/grizzly/main.py +++ b/grizzly/main.py @@ -129,7 +129,7 @@ def main(args): session.run( args.ignore, time_limit, - input_path=str(args.input) if args.input else None, + input_path=str(args.input.resolve()) if args.input else None, iteration_limit=args.limit, no_harness=args.no_harness, result_limit=1 if args.smoke_test else 0,