diff --git a/pathod/language/generators.py b/pathod/language/generators.py index 93db30148..1961df74d 100644 --- a/pathod/language/generators.py +++ b/pathod/language/generators.py @@ -1,7 +1,7 @@ +import os import string import random import mmap - import sys DATATYPES = dict( @@ -74,24 +74,20 @@ class RandomGenerator: class FileGenerator: - def __init__(self, path): self.path = path - self.fp = open(path, "rb") - self.map = mmap.mmap(self.fp.fileno(), 0, access=mmap.ACCESS_READ) def __len__(self): - return len(self.map) + return os.path.getsize(self.path) def __getitem__(self, x): - if isinstance(x, slice): - return self.map.__getitem__(x) - # A slice of length 1 returns a byte object (not an integer) - return self.map.__getitem__(slice(x, x + 1 or self.map.size())) + with open(self.path, mode="rb") as f: + if isinstance(x, slice): + with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mapped: + return mapped.__getitem__(x) + else: + f.seek(x) + return f.read(1) def __repr__(self): return "<%s" % self.path - - def close(self): - self.map.close() - self.fp.close() diff --git a/pathod/pathod_cmdline.py b/pathod/pathod_cmdline.py index ef1e983ff..dee19f4fc 100644 --- a/pathod/pathod_cmdline.py +++ b/pathod/pathod_cmdline.py @@ -216,7 +216,8 @@ def args_pathod(argv, stdout_=sys.stdout, stderr_=sys.stderr): anchors = [] for patt, spec in args.anchors: if os.path.isfile(spec): - data = open(spec).read() + with open(spec) as f: + data = f.read() spec = data try: arex = re.compile(patt) diff --git a/test/pathod/language/test_generators.py b/test/pathod/language/test_generators.py index 6a67ab72f..5e64c7266 100644 --- a/test/pathod/language/test_generators.py +++ b/test/pathod/language/test_generators.py @@ -14,16 +14,14 @@ def test_randomgenerator(): def test_filegenerator(tmpdir): f = tmpdir.join("foo") - f.write(b"x" * 10000) + f.write(b"abcdefghijklmnopqrstuvwxyz" * 1000) g = generators.FileGenerator(str(f)) - assert len(g) == 10000 - assert g[0] == b"x" - assert g[-1] == b"x" - assert g[0:5] == b"xxxxx" + assert len(g) == 26000 + assert g[0] == b"a" + assert g[2:7] == b"cdefg" assert len(g[1:10]) == 9 - assert len(g[10000:10001]) == 0 + assert len(g[26000:26001]) == 0 assert repr(g) - g.close() def test_transform_generator(): diff --git a/test/pathod/test_test.py b/test/pathod/test_test.py index 30d0a473e..d51a2c7a3 100644 --- a/test/pathod/test_test.py +++ b/test/pathod/test_test.py @@ -2,7 +2,6 @@ import os import requests import pytest -from mitmproxy.test import tutils from pathod import test from pathod.pathod import SSLOptions, CA_CERT_NAME @@ -27,7 +26,7 @@ class TestDaemonManual: d = test.Daemon(ssl=True, ssloptions=ssloptions) rsp = requests.get( "https://localhost:%s/p/202:da" % d.port, - verify=os.path.join(d.thread.server.ssloptions.confdir, CA_CERT_NAME)) + verify=os.path.expanduser(os.path.join(d.thread.server.ssloptions.confdir, CA_CERT_NAME))) assert rsp.ok assert rsp.status_code == 202 d.shutdown() diff --git a/test/pathod/tservers.py b/test/pathod/tservers.py index 3dc263117..a7c929646 100644 --- a/test/pathod/tservers.py +++ b/test/pathod/tservers.py @@ -1,3 +1,4 @@ +import os import tempfile import re import shutil