Optimize the handling of problematic files

2 files changed, 25 insertions, 1 deletions

diff --git a/tests/test_corrupted_files.py b/tests/test_corrupted_files.py
index 490ea42..82c6c3b 100644
--- a/tests/test_corrupted_files.py
+++ b/tests/test_corrupted_files.py
@@ -204,3 +204,14 @@ class TestCorruptedFiles(unittest.TestCase):
         p = video.AVIParser('./tests/data/clean.avi')
         self.assertFalse(p.remove_all())
         os.remove('./tests/data/clean.avi')
+
+    def test_avi_injection(self):
+        try:
+            video._get_ffmpeg_path()
+        except RuntimeError:
+            raise unittest.SkipTest
+
+        shutil.copy('./tests/data/dirty.torrent', './tests/data/--output.avi')
+        p = video.AVIParser('./tests/data/--output.avi')
+        self.assertFalse(p.remove_all())
+        os.remove('./tests/data/--output.avi')
diff --git a/tests/test_libmat2.py b/tests/test_libmat2.py
index 241c6eb..f5fc9e8 100644
--- a/tests/test_libmat2.py
+++ b/tests/test_libmat2.py
@@ -46,10 +46,23 @@ class TestParameterInjection(unittest.TestCase):
         shutil.copy('./tests/data/dirty.avi', './--output')
         p = video.AVIParser('--output')
         meta = p.get_meta()
-        print(meta)
         self.assertEqual(meta['Software'], 'MEncoder SVN-r33148-4.0.1')
         os.remove('--output')
 
+    def test_ffmpeg_injection_complete_path(self):
+        try:
+            video._get_ffmpeg_path()
+        except RuntimeError:
+            raise unittest.SkipTest
+
+        shutil.copy('./tests/data/dirty.avi', './tests/data/ --output.avi')
+        p = video.AVIParser('./tests/data/ --output.avi')
+        meta = p.get_meta()
+        self.assertEqual(meta['Software'], 'MEncoder SVN-r33148-4.0.1')
+        self.assertTrue(p.remove_all())
+        os.remove('./tests/data/ --output.avi')
+        os.remove('./tests/data/ --output.cleaned.avi')
+
 
 class TestUnsupportedEmbeddedFiles(unittest.TestCase):
     def test_odt_with_svg(self):