[loki] Use pathlib where possible

loki/args.py

@@ -2,6 +2,7 @@
 # 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 argparse import ArgumentParser
+from pathlib import Path
 
 from loki import Loki
 
@@ -10,7 +11,9 @@
 
 def parse_args(argv=None):
     parser = ArgumentParser(description="Loki fuzzing library")
-    parser.add_argument("input", help="Output will be generated based on this file")
+    parser.add_argument(
+        "input", type=Path, help="Output will be generated based on this file"
+    )
     parser.add_argument(
         "-a",
         "--aggression",
@@ -31,7 +34,13 @@ def parse_args(argv=None):
         "--count",
         default=1,
         type=int,
-        help="Number test cases to generate, minimum 1 (default: %(default)s)",
+        help="Number of test cases to generate, minimum 1 (default: %(default)s)",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        help="Output directory for fuzzed test cases (default: '.')",
     )
     parser.add_argument(
         "-q",
@@ -40,12 +49,9 @@ def parse_args(argv=None):
         action="store_true",
         help="Display limited output (default: %(default)s)",
     )
-    parser.add_argument(
-        "-o",
-        "--output",
-        default=None,
-        help="Output directory for fuzzed test cases (default: '.')",
-    )
     args = parser.parse_args(argv)
 
+    if not args.input.is_file():
+        parser.error(f"'{args.input}' is not a file")
+
     return args

loki/loki.py

@@ -5,15 +5,13 @@
 Loki fuzzing library
 """
 from logging import ERROR, INFO, basicConfig, getLogger
-from os import SEEK_END, makedirs
-from os.path import abspath, getsize
-from os.path import join as pathjoin
-from os.path import splitext
+from os import SEEK_END
+from pathlib import Path
 from random import choice, getrandbits, randint, sample
 from shutil import copy
 from struct import pack, unpack
 from tempfile import SpooledTemporaryFile, mkdtemp
-from time import strftime, time
+from time import perf_counter, strftime
 
 __author__ = "Tyson Smith"
 
@@ -50,7 +48,7 @@ def _fuzz_data(in_data, byte_order):
         elif fuzz_op == 1:  # arithmetic
             out_data = unpack(pack_unit, in_data)[0] + randint(-10, 10)
         elif fuzz_op == 2:  # interesting byte, short or int
-            out_data = choice((0, 1, int(mask / 2), int(mask / 2) + 1, mask))
+            out_data = choice((0, 1, mask // 2, (mask // 2) + 1, mask))
         elif fuzz_op == 3:  # random byte, short or int
             out_data = getrandbits(32)
         elif fuzz_op == 4:
@@ -103,7 +101,6 @@ def _fuzz(self, tgt_fp):
             tgt_fp.write(out_data)
 
     def fuzz_data(self, data):
-        assert data
         assert isinstance(data, bytes)
         # open a temp file in memory for fuzzing
         with SpooledTemporaryFile(max_size=0x800000, mode="r+b") as tmp_fp:
@@ -114,38 +111,38 @@ def fuzz_data(self, data):
 
     def fuzz_file(self, in_file, count, out_dir, ext=None):
         try:
-            if getsize(in_file) < 1:
+            if in_file.stat().st_size < 1:
                 LOG.error("Input must be at least 1 byte long")
                 return False
-        except OSError:
-            LOG.error("%r does not exists!", in_file)
+        except FileNotFoundError:
+            LOG.error("'%s' does not exist!", in_file)
             return False
         if ext is None:
-            ext = splitext(in_file)[1]
+            ext = in_file.suffix
         for i in range(count):
-            out_file = pathjoin(out_dir, f"{i:0>6d}_fuzzed{ext}")
+            out_file = out_dir / f"{i:06d}_fuzzed{ext}"
             copy(in_file, out_file)
-            with open(out_file, "r+b") as out_fp:
+            with out_file.open("r+b") as out_fp:
                 self._fuzz(out_fp)
         return True
 
     @classmethod
     def main(cls, args):
         basicConfig(format="", level=INFO if not args.quiet else ERROR)
         LOG.info("Starting Loki @ %s", strftime("%Y-%m-%d %H:%M:%S"))
-        LOG.info("Target template is %r", abspath(args.input))
+        LOG.info("Target template is '%s'", args.input.resolve())
         out_dir = args.output
         if out_dir is None:
-            out_dir = mkdtemp(prefix=strftime("loki_%m-%d_%H-%M_"), dir=".")
-        makedirs(out_dir, exist_ok=True)
-        LOG.info("Output directory is %r", abspath(out_dir))
+            out_dir = Path(mkdtemp(prefix=strftime("loki_%m-%d_%H-%M_"), dir="."))
+        out_dir.mkdir(exist_ok=True)
+        LOG.info("Output directory is '%s'", out_dir.resolve())
         count = max(args.count, 1)
         LOG.info("Generating %d fuzzed test cases...", count)
         loki = Loki(aggression=args.aggression, byte_order=args.byte_order)
         try:
-            start_time = time()
+            start_time = perf_counter()
             success = loki.fuzz_file(args.input, count, out_dir)
-            finish_time = time() - start_time
+            finish_time = perf_counter() - start_time
             LOG.info("Done. Total run time %gs", finish_time)
             if success:
                 LOG.info("About %gs per file", finish_time / count)

loki/test_loki.py

@@ -34,7 +34,7 @@ def test_loki_fuzz_file(tmp_path, in_size, aggression, byte_order):
 
     fuzzer = Loki(aggression=aggression, byte_order=byte_order)
     for _ in range(100):
-        assert fuzzer.fuzz_file(str(tmp_fn), 1, str(out_path))
+        assert fuzzer.fuzz_file(tmp_fn, 1, out_path)
         out_files = list(out_path.iterdir())
         assert len(out_files) == 1
         out_data = out_files[0].read_bytes()
@@ -49,14 +49,14 @@ def test_loki_01(tmp_path):
     """test Loki.fuzz_file() error cases"""
     fuzzer = Loki(aggression=0.1)
     # test missing file
-    assert not fuzzer.fuzz_file("nofile.test", 1, str(tmp_path))
+    assert not fuzzer.fuzz_file(tmp_path / "nofile.test", 1, tmp_path)
     assert not list(tmp_path.iterdir())
     # test empty file
     tmp_fn = tmp_path / "input"
     tmp_fn.touch()
     out_path = tmp_path / "out"
     out_path.mkdir()
-    assert not fuzzer.fuzz_file(str(tmp_fn), 1, str(out_path))
+    assert not fuzzer.fuzz_file(tmp_fn, 1, out_path)
     assert not list(out_path.iterdir())
 
 
@@ -158,12 +158,17 @@ def test_main_01(mocker, tmp_path):
     sample = tmp_path / "file.bin"
     sample.write_bytes(b"test!")
     args = mocker.Mock(
-        aggression=0.1, byte_order=None, count=15, input=str(sample), output=None
+        aggression=0.1, byte_order=None, count=15, input=sample, output=None
     )
     assert Loki.main(args) == 0
     assert fake_mkdtemp.call_count == 1
 
 
-def test_args_01():
+def test_args_01(capsys, tmp_path):
     """test parse_args()"""
-    assert parse_args(argv=["sample"])
+    with raises(SystemExit):
+        parse_args(argv=["missing.file"])
+    assert "error: 'missing.file' is not a file" in capsys.readouterr()[-1]
+    sample = tmp_path / "foo.txt"
+    sample.touch()
+    assert parse_args(argv=[str(sample)])