From 0bc67280679a460b285db76ec1243115e07e74a0 Mon Sep 17 00:00:00 2001 From: Christian Herdtweck Date: Mon, 29 Apr 2019 13:15:02 +0200 Subject: [PATCH] test: Use call_and_capture in olevba tests --- tests/olevba/test_basic.py | 126 +++++++++++++++++++++++++++++++++++++++--------------------------------------------------------------------------------------- 1 file changed, 39 insertions(+), 87 deletions(-) diff --git a/tests/olevba/test_basic.py b/tests/olevba/test_basic.py index abcfe8e..28238fc 100644 --- a/tests/olevba/test_basic.py +++ b/tests/olevba/test_basic.py @@ -3,44 +3,12 @@ Test basic functionality of olevba[3] """ import unittest -import sys import os from os.path import join -from contextlib import contextmanager -try: - from cStringIO import StringIO -except ImportError: # py3: - from io import StringIO -if sys.version_info.major <= 2: - from oletools import olevba -else: - from oletools import olevba3 as olevba +import re # Directory with test data, independent of current working directory -from tests.test_utils import DATA_BASE_DIR - - -@contextmanager -def capture_output(): - """ - Temporarily replace stdout/stderr with buffers to capture output. - - Once we only support python>=3.4: this is already built into python as - :py:func:`contextlib.redirect_stdout`. - - Not quite sure why, but seems to only work once per test function ... - """ - orig_stdout = sys.stdout - orig_stderr = sys.stderr - - try: - sys.stdout = StringIO() - sys.stderr = StringIO() - yield sys.stdout, sys.stderr - - finally: - sys.stdout = orig_stdout - sys.stderr = orig_stderr +from tests.test_utils import DATA_BASE_DIR, call_and_capture class TestOlevbaBasic(unittest.TestCase): @@ -57,62 +25,48 @@ class TestOlevbaBasic(unittest.TestCase): def do_test_behaviour(self, filename): """Helper for test_{text,empty}_behaviour.""" input_file = join(DATA_BASE_DIR, 'basic', filename) - ret_code = -1 - - # run olevba, capturing its output and return code - with capture_output() as (stdout, stderr): - with self.assertRaises(SystemExit) as raise_context: - olevba.main([input_file, ]) - ret_code = raise_context.exception.code - - # check that return code is 0 - self.assertEqual(ret_code, 0) - - # check there are only warnings in stderr - stderr = stderr.getvalue() - skip_line = False - for line in stderr.splitlines(): - if skip_line: - skip_line = False - continue - self.assertTrue(line.startswith('WARNING ') or - 'ResourceWarning' in line, - msg='Line "{}" in stderr is unexpected for {}'\ - .format(line.rstrip(), filename)) - if 'ResourceWarning' in line: - skip_line = True - self.assertIn('not encrypted', stderr) - - # check stdout - stdout = stdout.getvalue().lower() - self.assertIn(input_file.lower(), stdout) - self.assertIn('type: text', stdout) - self.assertIn('no suspicious', stdout) - self.assertNotIn('error', stdout) - self.assertNotIn('warn', stdout) + output, _ = call_and_capture('olevba', args=(input_file, )) + + # check output + self.assertTrue(re.search(r'^Type:\s+Text\s*$', output, re.MULTILINE), + msg='"Type: Text" not found in output:\n' + output) + self.assertTrue(re.search(r'^No suspicious .+ found.$', output, + re.MULTILINE), + msg='"No suspicous...found" not found in output:\n' + \ + output) + self.assertNotIn('error', output.lower()) + + # check warnings + for line in output.splitlines(): + if line.startswith('WARNING ') and 'encrypted' in line: + continue # encryption warnings are ok + elif 'warn' in line.lower(): + raise self.fail('Found "warn" in output line: "{}"' + .format(line.rstrip())) + self.assertIn('not encrypted', output) def test_rtf_behaviour(self): """Test behaviour of olevba when presented with an rtf file.""" input_file = join(DATA_BASE_DIR, 'msodde', 'RTF-Spec-1.7.rtf') - ret_code = -1 - - # run olevba, capturing its output and return code - with capture_output() as (stdout, stderr): - with self.assertRaises(SystemExit) as raise_context: - olevba.main([input_file, ]) - ret_code = raise_context.exception.code + output, ret_code = call_and_capture('olevba', args=(input_file, ), + accept_nonzero_exit=True) # check that return code is olevba.RETURN_OPEN_ERROR self.assertEqual(ret_code, 5) - stdout = stdout.getvalue().lower() - self.assertNotIn('error', stdout) - self.assertNotIn('warn', stdout) - stderr = stderr.getvalue().lower() - self.assertIn('fileopenerror', stderr) - self.assertIn('is rtf', stderr) - self.assertIn('rtfobj.py', stderr) - self.assertIn('not encrypted', stderr) + # check output: + self.assertIn('FileOpenError', output) + self.assertIn('is RTF', output) + self.assertIn('rtfobj.py', output) + self.assertIn('not encrypted', output) + + # check warnings + for line in output.splitlines(): + if line.startswith('WARNING ') and 'encrypted' in line: + continue # encryption warnings are ok + elif 'warn' in line.lower(): + raise self.fail('Found "warn" in output line: "{}"' + .format(line.rstrip())) def test_crypt_return(self): """ @@ -136,11 +90,9 @@ class TestOlevbaBasic(unittest.TestCase): continue full_name = join(CRYPT_DIR, filename) for args in ADD_ARGS: - try: - olevba.main(args + [full_name, ]) - self.fail('Olevba should have exited') - except SystemExit as sys_exit: - ret_code = sys_exit.code or 0 # sys_exit.code can be None + _, ret_code = call_and_capture('olevba', + args=[full_name, ] + args, + accept_nonzero_exit=True) self.assertEqual(ret_code, CRYPT_RETURN_CODE, msg='Wrong return code {} for args {}'\ .format(ret_code, args + [filename, ])) -- libgit2 0.21.4