diff options
| -rwxr-xr-x | mat2 | 16 | ||||
| -rw-r--r-- | tests/test_climat2.py | 6 |
2 files changed, 17 insertions, 5 deletions
| @@ -41,6 +41,9 @@ def create_arg_parser(): | |||
| 41 | help='check if MAT2 has all the dependencies it needs') | 41 | help='check if MAT2 has all the dependencies it needs') |
| 42 | parser.add_argument('-V', '--verbose', action='store_true', | 42 | parser.add_argument('-V', '--verbose', action='store_true', |
| 43 | help='show more verbose status information') | 43 | help='show more verbose status information') |
| 44 | parser.add_argument('-u', '--unknown-members', metavar='POLICY', default='abort', | ||
| 45 | help='how to handle unknown members of archive-style files ' + | ||
| 46 | '(POLICY should be abort, omit, or keep)') | ||
| 44 | 47 | ||
| 45 | 48 | ||
| 46 | info = parser.add_mutually_exclusive_group() | 49 | info = parser.add_mutually_exclusive_group() |
| @@ -67,8 +70,8 @@ def show_meta(filename: str): | |||
| 67 | except UnicodeEncodeError: | 70 | except UnicodeEncodeError: |
| 68 | print(" %s: harmful content" % k) | 71 | print(" %s: harmful content" % k) |
| 69 | 72 | ||
| 70 | def clean_meta(params: Tuple[str, bool]) -> bool: | 73 | def clean_meta(params: Tuple[str, bool, str]) -> bool: |
| 71 | filename, is_lightweight = params | 74 | filename, is_lightweight, unknown_member_policy = params |
| 72 | if not __check_file(filename, os.R_OK|os.W_OK): | 75 | if not __check_file(filename, os.R_OK|os.W_OK): |
| 73 | return False | 76 | return False |
| 74 | 77 | ||
| @@ -76,6 +79,7 @@ def clean_meta(params: Tuple[str, bool]) -> bool: | |||
| 76 | if p is None: | 79 | if p is None: |
| 77 | print("[-] %s's format (%s) is not supported" % (filename, mtype)) | 80 | print("[-] %s's format (%s) is not supported" % (filename, mtype)) |
| 78 | return False | 81 | return False |
| 82 | p.unknown_member_policy = unknown_member_policy | ||
| 79 | if is_lightweight: | 83 | if is_lightweight: |
| 80 | return p.remove_all_lightweight() | 84 | return p.remove_all_lightweight() |
| 81 | return p.remove_all() | 85 | return p.remove_all() |
| @@ -133,9 +137,15 @@ def main(): | |||
| 133 | return 0 | 137 | return 0 |
| 134 | 138 | ||
| 135 | else: | 139 | else: |
| 140 | if args.unknown_members == 'keep': | ||
| 141 | logging.warning('Keeping unknown member files may leak metadata in the resulting file!') | ||
| 142 | elif args.unknown_members not in ['omit', 'abort']: | ||
| 143 | logging.warning('Undefined policy for handling unknown member files: "%s"', | ||
| 144 | args.unknown_members) | ||
| 136 | p = multiprocessing.Pool() | 145 | p = multiprocessing.Pool() |
| 137 | mode = (args.lightweight is True) | 146 | mode = (args.lightweight is True) |
| 138 | l = zip(__get_files_recursively(args.files), itertools.repeat(mode)) | 147 | l = zip(__get_files_recursively(args.files), itertools.repeat(mode), |
| 148 | itertools.repeat(args.unknown_members)) | ||
| 139 | 149 | ||
| 140 | ret = list(p.imap_unordered(clean_meta, list(l))) | 150 | ret = list(p.imap_unordered(clean_meta, list(l))) |
| 141 | return 0 if all(ret) else -1 | 151 | return 0 if all(ret) else -1 |
diff --git a/tests/test_climat2.py b/tests/test_climat2.py index af89c0e..6ee84d5 100644 --- a/tests/test_climat2.py +++ b/tests/test_climat2.py | |||
| @@ -8,12 +8,14 @@ class TestHelp(unittest.TestCase): | |||
| 8 | def test_help(self): | 8 | def test_help(self): |
| 9 | proc = subprocess.Popen(['./mat2', '--help'], stdout=subprocess.PIPE) | 9 | proc = subprocess.Popen(['./mat2', '--help'], stdout=subprocess.PIPE) |
| 10 | stdout, _ = proc.communicate() | 10 | stdout, _ = proc.communicate() |
| 11 | self.assertIn(b'usage: mat2 [-h] [-v] [-l] [-c] [-V] [-s | -L] [files [files ...]]', stdout) | 11 | self.assertIn(b'usage: mat2 [-h] [-v] [-l] [-c] [-V] [-u POLICY] [-s | -L] [files [files ...]]', |
| 12 | stdout) | ||
| 12 | 13 | ||
| 13 | def test_no_arg(self): | 14 | def test_no_arg(self): |
| 14 | proc = subprocess.Popen(['./mat2'], stdout=subprocess.PIPE) | 15 | proc = subprocess.Popen(['./mat2'], stdout=subprocess.PIPE) |
| 15 | stdout, _ = proc.communicate() | 16 | stdout, _ = proc.communicate() |
| 16 | self.assertIn(b'usage: mat2 [-h] [-v] [-l] [-c] [-V] [-s | -L] [files [files ...]]', stdout) | 17 | self.assertIn(b'usage: mat2 [-h] [-v] [-l] [-c] [-V] [-u POLICY] [-s | -L] [files [files ...]]', |
| 18 | stdout) | ||
| 17 | 19 | ||
| 18 | 20 | ||
| 19 | class TestVersion(unittest.TestCase): | 21 | class TestVersion(unittest.TestCase): |
