Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
55 changes: 28 additions & 27 deletions MitmLibrary/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@
applications in a more realistic and controlled environment.
"""

from typing import Any, Dict, List, Optional, Sequence, Union
from collections.abc import Sequence
from typing import Any

from mitmproxy.tools import dump
from robot.api import logger
Expand Down Expand Up @@ -202,15 +203,15 @@ def __init__(self) -> None:
"""
self.controller: ProxyController = ProxyController()
self.registry: RuleRegistry = RuleRegistry()
self.interceptor: Optional[Interceptor] = None
self.recorder: Optional[FlowRecorder] = None
self.interceptor: Interceptor | None = None
self.recorder: FlowRecorder | None = None
self.log_to_console: bool = True
# Robot Framework calls close() on this when the suite that imported the library
# ends, which releases the port even if the suite never stopped the proxy itself.
self.ROBOT_LIBRARY_LISTENER: LibraryListener = LibraryListener(self.controller)

@property
def proxy_master(self) -> Optional[dump.DumpMaster]:
def proxy_master(self) -> dump.DumpMaster | None:
"""The running mitmproxy master, or None when no proxy is running."""
return self.controller.master

Expand Down Expand Up @@ -239,14 +240,14 @@ def start_mitm_proxy(
self,
listen_host: str = "127.0.0.1",
listen_port: int = 8080,
certificates_directory: Optional[str] = None,
certificates_directory: str | None = None,
ssl_insecure: bool = False,
log_to_console: bool = True,
record: bool = False,
record_limit: int = DEFAULT_LIMIT,
record_body_limit: int = DEFAULT_BODY_LIMIT,
mode: Optional[Union[str, List[str]]] = None,
proxy_auth: Optional[str] = None,
mode: str | list[str] | None = None,
proxy_auth: str | None = None,
) -> None:
"""
Starts a proxy at the given host and port.
Expand Down Expand Up @@ -309,7 +310,7 @@ def _build_addons(self, master: dump.DumpMaster) -> Sequence[Any]:
there afterwards. Only the addon reading them is rebuilt.
"""
self.interceptor = Interceptor(self.registry, self.log_to_console)
addons: List[Any] = [self.interceptor]
addons: list[Any] = [self.interceptor]
if self.recorder is not None:
addons.append(self.recorder)
return addons
Expand Down Expand Up @@ -362,7 +363,7 @@ def block_requests(
url: str,
mode: BlockMode = BlockMode.RESPOND,
status_code: int = 403,
body: Optional[str] = None,
body: str | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
times: int = 0,
Expand Down Expand Up @@ -397,8 +398,8 @@ def set_response(
alias: str,
url: str,
status_code: int = 200,
headers: Optional[Dict[str, str]] = None,
body: Optional[str] = None,
headers: dict[str, str] | None = None,
body: str | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
times: int = 0,
Expand Down Expand Up @@ -510,8 +511,8 @@ def set_response_headers(
self,
alias: str,
url: str,
headers: Optional[Dict[str, str]] = None,
remove: Optional[List[str]] = None,
headers: dict[str, str] | None = None,
remove: list[str] | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
times: int = 0,
Expand Down Expand Up @@ -574,8 +575,8 @@ def set_request_headers(
self,
alias: str,
url: str,
headers: Optional[Dict[str, str]] = None,
remove: Optional[List[str]] = None,
headers: dict[str, str] | None = None,
remove: list[str] | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
times: int = 0,
Expand Down Expand Up @@ -671,8 +672,8 @@ def redirect_requests_to_host(
alias: str,
url: str,
host: str,
port: Optional[int] = None,
scheme: Optional[str] = None,
port: int | None = None,
scheme: str | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
times: int = 0,
Expand Down Expand Up @@ -743,7 +744,7 @@ def simulate_truncated_response(
self,
alias: str,
url: str,
keep_bytes: Optional[int] = None,
keep_bytes: int | None = None,
keep_fraction: float = 0.5,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
Expand Down Expand Up @@ -806,7 +807,7 @@ def clear_all_rules(self) -> None:
self._require_registry().clear()

@keyword
def get_proxy_rules(self) -> List[DotDict]:
def get_proxy_rules(self) -> list[DotDict]:
"""Returns the loaded rules, in the order they are applied.

Each rule is a dictionary with at least `alias`, `url`, `match`, `method`,
Expand Down Expand Up @@ -879,10 +880,10 @@ def clear_recorded_requests(self) -> None:
@keyword
def get_recorded_requests(
self,
url: Optional[str] = None,
url: str | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
) -> List[DotDict]:
) -> list[DotDict]:
"""Returns the recorded requests, oldest first.

Each request is a dictionary with `method`, `url`, `host`, `path`, `query`,
Expand All @@ -904,7 +905,7 @@ def get_recorded_requests(
@keyword
def get_request_count(
self,
url: Optional[str] = None,
url: str | None = None,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
) -> int:
Expand All @@ -928,8 +929,8 @@ def request_should_have_been_made(
url: str,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
times: Optional[int] = None,
msg: Optional[str] = None,
times: int | None = None,
msg: str | None = None,
) -> None:
"""Fails unless a matching request was recorded.

Expand Down Expand Up @@ -972,7 +973,7 @@ def request_should_not_have_been_made(
url: str,
method: str = ANY_METHOD,
match: MatchMode = MatchMode.SUBSTRING,
msg: Optional[str] = None,
msg: str | None = None,
) -> None:
"""Fails when a matching request was recorded.

Expand Down Expand Up @@ -1008,7 +1009,7 @@ def wait_until_request_is_made(
match: MatchMode = MatchMode.SUBSTRING,
timeout: str = "10s",
count: int = 1,
) -> List[DotDict]:
) -> list[DotDict]:
"""Waits until matching requests have been recorded, and returns them.

For traffic a test does not trigger directly, such as a request a page makes
Expand Down Expand Up @@ -1041,7 +1042,7 @@ def _require_recorder(self) -> FlowRecorder:

@not_keyword
def _recording_matcher(
self, url: Optional[str], method: str, match: MatchMode
self, url: str | None, method: str, match: MatchMode
) -> UrlMatcher:
"""Builds the matcher the recording keywords filter with.

Expand Down
4 changes: 2 additions & 2 deletions MitmLibrary/async_loop_thread.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ def run(self) -> None:
asyncio.set_event_loop(self.loop)
try:
self.loop.run_forever()
except Exception as e: # pylint: disable=broad-exception-caught
except Exception as e: # noqa: BLE001 - best-effort loop shutdown, must not raise
print(f"Async loop thread error: {e}") # Log the error message

def stop(self, timeout: float = 5) -> None:
Expand All @@ -51,5 +51,5 @@ def stop(self, timeout: float = 5) -> None:
self.join(timeout=timeout)
try:
self.loop.close()
except Exception as error: # pylint: disable=broad-exception-caught
except Exception as error: # noqa: BLE001 - best-effort shutdown, must not raise
print(f"Async loop thread could not be closed: {error}")
8 changes: 4 additions & 4 deletions MitmLibrary/failures.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@

import asyncio
from dataclasses import dataclass
from typing import Any, Dict, Optional
from typing import Any

from mitmproxy import http
from robot.api import logger
Expand Down Expand Up @@ -46,7 +46,7 @@ async def apply_async(self, flow: http.HTTPFlow) -> bool:
kill_flow(flow)
return True

def describe(self) -> Dict[str, Any]:
def describe(self) -> dict[str, Any]:
return {"type": "timeout", "hold": self.hold, "hold_seconds": self.hold_seconds}


Expand All @@ -61,7 +61,7 @@ class TruncateAction(Action):
with the response at all.
"""

keep_bytes: Optional[int] = None
keep_bytes: int | None = None
keep_fraction: float = 0.5

phase = Phase.RESPONSE
Expand Down Expand Up @@ -92,7 +92,7 @@ def _keep(self, length: int) -> int:
return max(0, self.keep_bytes)
return max(0, int(length * self.keep_fraction))

def describe(self) -> Dict[str, Any]:
def describe(self) -> dict[str, Any]:
return {
"type": "truncate",
"keep_bytes": self.keep_bytes,
Expand Down
3 changes: 1 addition & 2 deletions MitmLibrary/interceptor.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
only decides *when* they run, not *what* they do.
"""

from typing import Optional

from mitmproxy import http
from robot.api import logger
Expand Down Expand Up @@ -69,7 +68,7 @@ def _log(self, rule: Rule, flow: http.HTTPFlow) -> None:
)


def _method(flow: http.HTTPFlow) -> Optional[str]:
def _method(flow: http.HTTPFlow) -> str | None:
"""The request method, or None when the flow does not report one."""
method = getattr(flow.request, "method", None)
return method if isinstance(method, str) else None
10 changes: 5 additions & 5 deletions MitmLibrary/matching.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
import re
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, Pattern
from re import Pattern

ANY_METHOD = "ANY"

Expand Down Expand Up @@ -47,7 +47,7 @@ class UrlMatcher:
pattern: str
mode: MatchMode = MatchMode.SUBSTRING
method: str = ANY_METHOD
_regex: Optional[Pattern[str]] = field(default=None, init=False, repr=False)
_regex: Pattern[str] | None = field(default=None, init=False, repr=False)

def __post_init__(self) -> None:
# The dataclass is frozen so that a rule cannot change what it matches while the
Expand All @@ -56,7 +56,7 @@ def __post_init__(self) -> None:
object.__setattr__(self, "method", (self.method or ANY_METHOD).strip().upper())
object.__setattr__(self, "_regex", self._compile())

def _compile(self) -> Optional[Pattern[str]]:
def _compile(self) -> Pattern[str] | None:
"""Builds the expression for the mode, or None when the mode needs no expression."""
if self.mode is MatchMode.SUBSTRING:
return None
Expand All @@ -77,13 +77,13 @@ def matches_url(self, url: str) -> bool:
return self._regex.match(url) is not None
return self._regex.search(url) is not None

def matches_method(self, method: Optional[str]) -> bool:
def matches_method(self, method: str | None) -> bool:
"""Whether the method matches. `ANY` matches everything, including no method."""
if self.method == ANY_METHOD:
return True
return (method or "").upper() == self.method

def matches(self, url: str, method: Optional[str] = None) -> bool:
def matches(self, url: str, method: str | None = None) -> bool:
"""Whether both the url and the method match."""
return self.matches_method(method) and self.matches_url(url)

Expand Down
Loading