maplocal addon: minor improvements

This commit is contained in:
Martin Plattner 2020-07-08 01:34:14 +02:00
parent 7022a54737
commit b1609697cd
2 changed files with 14 additions and 18 deletions

View File

@ -41,15 +41,15 @@ class MapLocal:
def sanitize_candidate_path(self, candidate_path, base_path): def sanitize_candidate_path(self, candidate_path, base_path):
try: try:
candidate_path.resolve(strict=True) candidate_path = candidate_path.resolve(strict=True)
if base_path in candidate_path.parents: if base_path == candidate_path or base_path in candidate_path.parents:
return candidate_path return candidate_path
except FileNotFoundError: except FileNotFoundError:
pass pass
return None return None
def file_candidates(self, url: str, spec: ModifySpec) -> typing.List[Path]: def file_candidates(self, url: str, spec: ModifySpec) -> typing.List[Path]:
replacement = spec.replacement_str replacement = spec.replacement
candidates = [] candidates = []
if replacement.is_file(): if replacement.is_file():
@ -85,14 +85,13 @@ class MapLocal:
) )
return mimetype return mimetype
def request(self, flow: http.HTTPFlow) -> None: def request(self, flow: http.HTTPFlow) -> None:
if flow.reply and flow.reply.has_message: if flow.reply and flow.reply.has_message:
return return
for spec in self.replacements: for spec in self.replacements:
req = flow.request req = flow.request
url = req.pretty_url url = req.pretty_url
base_path = Path(spec.replacement_str) base_path = Path(spec.replacement)
if spec.matches(flow) and re.search(spec.subject, url.encode("utf8", "surrogateescape")): if spec.matches(flow) and re.search(spec.subject, url.encode("utf8", "surrogateescape")):
file_candidates = self.file_candidates(url, spec) file_candidates = self.file_candidates(url, spec)
@ -101,14 +100,11 @@ class MapLocal:
file_candidate = Path(file_candidate) file_candidate = Path(file_candidate)
if self.sanitize_candidate_path(file_candidate, base_path): if self.sanitize_candidate_path(file_candidate, base_path):
try: try:
with open(file_candidate, "rb") as file: flow.response = http.HTTPResponse.make(
replacement = file.read() 200,
file_candidate.read_bytes(),
{"Content-Type": self.get_mime_type(str(file_candidate))}
)
except IOError: except IOError:
ctx.log.warn(f"Could not read replacement file {file_candidate}") ctx.log.warn(f"Could not read replacement file {file_candidate}")
return return
flow.response = http.HTTPResponse.make(
200,
replacement,
{"Content-Type": self.get_mime_type(str(file_candidate))}
)

View File

@ -12,7 +12,7 @@ from mitmproxy import ctx
class ModifySpec(typing.NamedTuple): class ModifySpec(typing.NamedTuple):
matches: flowfilter.TFilter matches: flowfilter.TFilter
subject: bytes subject: bytes
replacement_str: str replacement: str
def read_replacement(self) -> bytes: def read_replacement(self) -> bytes:
""" """
@ -22,11 +22,11 @@ class ModifySpec(typing.NamedTuple):
Raises: Raises:
- IOError if the file cannot be read. - IOError if the file cannot be read.
""" """
if self.replacement_str.startswith("@"): if self.replacement.startswith("@"):
return Path(self.replacement_str[1:]).expanduser().read_bytes() return Path(self.replacement[1:]).expanduser().read_bytes()
else: else:
# We could cache this at some point, but unlikely to be a problem. # We could cache this at some point, but unlikely to be a problem.
return strutils.escaped_str_to_bytes(self.replacement_str) return strutils.escaped_str_to_bytes(self.replacement)
def _match_all(flow) -> bool: def _match_all(flow) -> bool: