Skip to content
Snippets Groups Projects
Commit d2b2a54a authored by Julien (jvoisin) Voisin's avatar Julien (jvoisin) Voisin
Browse files

MAT2's cli now uses meaningful return codes

- Simplify the multiprocessing by using a Pool
- Use some functional (:hearts:) constructions to exit
  with a return code
- Add some tests to prove that we're doing things
  that are working correctly
parent a79c9410
No related branches found
No related tags found
No related merge requests found
#!/usr/bin/python3 #!/usr/bin/python3
import os import os
from typing import Tuple
import sys
import itertools
import mimetypes import mimetypes
import argparse import argparse
from threading import Thread
import multiprocessing import multiprocessing
from queue import Queue
from src import parser_factory from src import parser_factory
...@@ -52,14 +53,15 @@ def show_meta(filename:str): ...@@ -52,14 +53,15 @@ def show_meta(filename:str):
print(" %s: harmful content" % k) print(" %s: harmful content" % k)
def clean_meta(filename:str, is_lightweigth:bool) -> bool: def clean_meta(params:Tuple[str, bool]) -> bool:
filename, is_lightweigth = params
if not __check_file(filename, os.R_OK|os.W_OK): if not __check_file(filename, os.R_OK|os.W_OK):
return return
p, mtype = parser_factory.get_parser(filename) p, mtype = parser_factory.get_parser(filename)
if p is None: if p is None:
print("[-] %s's format (%s) is not supported" % (filename, mtype)) print("[-] %s's format (%s) is not supported" % (filename, mtype))
return return False
if is_lightweigth: if is_lightweigth:
return p.remove_all_lightweight() return p.remove_all_lightweight()
return p.remove_all() return p.remove_all()
...@@ -82,15 +84,6 @@ def __get_files_recursively(files): ...@@ -82,15 +84,6 @@ def __get_files_recursively(files):
for _f in _files: for _f in _files:
yield os.path.join(path, _f) yield os.path.join(path, _f)
def __do_clean_async(is_lightweigth, q):
while True:
f = q.get()
if f is None: # nothing more to process
return
clean_meta(f, is_lightweigth)
q.task_done()
def main(): def main():
arg_parser = create_arg_parser() arg_parser = create_arg_parser()
args = arg_parser.parse_args() args = arg_parser.parse_args()
...@@ -106,24 +99,13 @@ def main(): ...@@ -106,24 +99,13 @@ def main():
show_meta(f) show_meta(f)
return return
else: # Thread the cleaning else:
p = multiprocessing.Pool()
mode = (args.lightweight is True) mode = (args.lightweight is True)
q = Queue(maxsize=0) l = zip(__get_files_recursively(args.files), itertools.repeat(mode))
threads = list()
for f in __get_files_recursively(args.files):
q.put(f)
for _ in range(multiprocessing.cpu_count()):
worker = Thread(target=__do_clean_async, args=(mode, q))
worker.start()
threads.append(worker)
for _ in range(multiprocessing.cpu_count()):
q.put(None) # stop the threads
for worker in threads:
worker.join()
ret = list(p.imap_unordered(clean_meta, list(l)))
return 0 if all(ret) else -1
if __name__ == '__main__': if __name__ == '__main__':
main() sys.exit(main())
...@@ -16,6 +16,22 @@ class TestHelp(unittest.TestCase): ...@@ -16,6 +16,22 @@ class TestHelp(unittest.TestCase):
self.assertIn(b'usage: main.py [-h] [-c] [-l] [-s] [-L] [files [files ...]]', stdout) self.assertIn(b'usage: main.py [-h] [-c] [-l] [-s] [-L] [files [files ...]]', stdout)
class TestReturnValue(unittest.TestCase):
def test_nonzero(self):
ret = subprocess.call(['./main.py', './main.py'], stdout=subprocess.DEVNULL)
self.assertEqual(255, ret)
ret = subprocess.call(['./main.py', '--whololo'], stderr=subprocess.DEVNULL)
self.assertEqual(2, ret)
def test_zero(self):
ret = subprocess.call(['./main.py'], stdout=subprocess.DEVNULL)
self.assertEqual(0, ret)
ret = subprocess.call(['./main.py', '--show', './main.py'], stdout=subprocess.DEVNULL)
self.assertEqual(0, ret)
class TestCleanMeta(unittest.TestCase): class TestCleanMeta(unittest.TestCase):
def test_jpg(self): def test_jpg(self):
shutil.copy('./tests/data/dirty.jpg', './tests/data/clean.jpg') shutil.copy('./tests/data/dirty.jpg', './tests/data/clean.jpg')
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment