diff --git a/loki/args.py b/loki/args.py index d948329d..d2de9507 100644 --- a/loki/args.py +++ b/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 diff --git a/loki/loki.py b/loki/loki.py index 1d4af843..537f3de1 100644 --- a/loki/loki.py +++ b/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,18 +111,18 @@ 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 @@ -133,19 +130,19 @@ def fuzz_file(self, in_file, count, out_dir, ext=None): 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) diff --git a/loki/test_loki.py b/loki/test_loki.py index fd3dd521..9e957a7b 100644 --- a/loki/test_loki.py +++ b/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)])