replaced engine by class and did some further restructuring

pdf_generator/generator.py

+import logging
+import os
+import shutil
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from common_helper_process import execute_shell_command_get_return_code
+from pdf_generator.pre_processing.rest import create_request_url, request_firmware_data
+from pdf_generator.tex_generation.template_engine import Engine
+
+
+def execute_latex(tmp_dir):
+    current_dir = os.getcwd()
+    os.chdir(tmp_dir)
+    execute_shell_command_get_return_code('env buf_size=1000000 pdflatex main.tex')
+    os.chdir(current_dir)
+
+
+def copy_fact_image(target):
+    shutil.copy(str(Path(__file__).parent / 'templates' / 'fact_logo.png'), str(Path(target) / 'fact_logo.png'))
+
+
+def generate_analysis_codes(engine, analysis):
+    return [
+        ('{}.tex'.format(analysis_plugin), engine.render_analysis_template(analysis_plugin, analysis[analysis_plugin])) for analysis_plugin in analysis
+    ]
+
+
+def create_report_filename(meta_data):
+    return '{}_analysis_report.pdf'.format(meta_data['device_name'].replace(' ', '_').replace('/', '__'))
+
+
+def generate_report(firmware_uid):
+    request_url = create_request_url(firmware_uid)
+    try:
+        analysis, meta_data = request_firmware_data(request_url)
+    except KeyError:
+        logging.warning('No firmware found with UID {}'.format(firmware_uid))
+        return 1
+
+    with TemporaryDirectory() as tmp_dir:
+        engine = Engine(tmp_dir=tmp_dir)
+        Path(tmp_dir, 'meta.tex').write_text(engine.render_meta_template(meta_data))
+
+        for filename, result_code in generate_analysis_codes(engine=engine, analysis=analysis):
+            Path(tmp_dir, filename).write_text(result_code)
+
+        Path(tmp_dir, 'main.tex').write_text(engine.render_main_template(analysis=analysis, meta_data=meta_data))
+
+        copy_fact_image(tmp_dir)
+
+        execute_latex(tmp_dir)
+
+        pdf_filename = create_report_filename(meta_data)
+        shutil.move(str(Path(tmp_dir, 'main.pdf')), str(Path('.', pdf_filename)))
+
+    return 0

pdf_generator/tex_generation/code_generation.py

-import logging
-
-from jinja2 import TemplateNotFound
-
-GENERIC_TEMPLATE = 'generic.tex'
-
-
-def generate_meta_data_code(environment, meta_data):
-    template = environment.get_template('meta_data.tex')
-    logging.debug('Rendering meta data')
-    return template.render(meta_data=meta_data)
-
-
-def generate_main_code(firmware_analyses, firmware_meta_data, jinja_environment):
-    template = jinja_environment.get_template('main.tex')
-    logging.debug('Rendering main page')
-    return template.render(analysis=firmware_analyses, meta_data=firmware_meta_data)
-
-
-def generate_analysis_codes(environment, analysis, tmp_dir):
-    return [
-        ('{}.tex'.format(analysis_plugin), _render_analysis_result(analysis[analysis_plugin], environment, analysis_plugin, tmp_dir)) for analysis_plugin in analysis
-    ]
-
-
-def _render_analysis_result(analysis, environment, analysis_plugin, tmp_dir):
-    try:
-        template = environment.get_template('{}.tex'.format(analysis_plugin))
-    except TemplateNotFound:
-        logging.debug('Falling back on generic template for {}'.format(analysis_plugin))
-        template = environment.get_template(GENERIC_TEMPLATE)
-
-    logging.debug('Rendering {}'.format(analysis_plugin))
-    return template.render(selected_analysis=analysis, tmp_dir=tmp_dir)
-
-
-def create_report_filename(meta_data):
-    main_tex_filename = meta_data['device_name'] + '_analysis_report.pdf'
-    main_tex_filename = main_tex_filename.replace(' ', '_')
-    return main_tex_filename.replace('/', '__')

pdf_generator/tex_generation/template_engine.py

+import logging
 from base64 import decodebytes
 from pathlib import Path
 from time import localtime, strftime
@@ -5,9 +6,11 @@ from time import localtime, strftime
 import jinja2
 from common_helper_files import human_readable_file_size
 
+GENERIC_TEMPLATE = 'generic.tex'
+
 
 def byte_number_filter(number, verbose=True):
-    if isinstance(number, int) or isinstance(number, float):
+    if isinstance(number, (int, float)):
         if verbose:
             return '{} ({})'.format(human_readable_file_size(int(number)), format(number, ',d') + ' bytes')
         return human_readable_file_size(int(number))
@@ -19,25 +22,23 @@ def nice_unix_time(unix_time_stamp):
     input unix_time_stamp
     output string 'YYYY-MM-DD HH:MM:SS'
     '''
-    if isinstance(unix_time_stamp, float) or isinstance(unix_time_stamp, int):
+    if isinstance(unix_time_stamp, (int, float)):
         tmp = localtime(unix_time_stamp)
         return strftime('%Y-%m-%d %H:%M:%S', tmp)
-    else:
-        return unix_time_stamp
+    return unix_time_stamp
 
 
-def nice_number_filter(i):
-    if isinstance(i, int):
-        return '{:,}'.format(i)
-    elif isinstance(i, float):
-        return '{:,.2f}'.format(i)
-    elif i is None:
+def nice_number_filter(number):
+    if isinstance(number, int):
+        return '{:,}'.format(number)
+    if isinstance(number, float):
+        return '{:,.2f}'.format(number)
+    if number is None:
         return 'not available'
-    else:
-        return i
+    return number
 
 
-def filter_latex_special_chars(data):
+def filter_latex_special_chars(data):  # pylint: disable=too-complex,too-many-branches
     if '\\' in data:
         data = data.replace('\\', '')
     if '$' in data:
@@ -87,12 +88,8 @@ def convert_base64_to_png_filter(base64_string, filename, directory):
     return str(file_path)
 
 
-def check_if_list_empty(ls):
-    if ls:
-        return ls
-    else:
-        empty_ls = ['list is empty']
-        return empty_ls
+def check_if_list_empty(list_of_strings):
+    return list_of_strings if list_of_strings else ['list is empty']
 
 
 def filter_chars_in_list(list_of_strings):
@@ -150,3 +147,25 @@ def _add_filters_to_jinja(environment):
     environment.filters['filter_list'] = filter_chars_in_list
     environment.filters['split_hash'] = split_hash
     environment.filters['split_output_lines'] = split_output_lines
+
+
+class Engine:
+    def __init__(self, template_folder=None, tmp_dir=None):
+        self._environment = create_jinja_environment(template_folder if template_folder else 'default')
+        self._tmp_dir = tmp_dir
+
+    def render_main_template(self, analysis, meta_data):
+        template = self._environment.get_template('main.tex')
+        return template.render(analysis=analysis, meta_data=meta_data)
+
+    def render_meta_template(self, meta_data):
+        template = self._environment.get_template('meta_data.tex')
+        return template.render(meta_data=meta_data)
+
+    def render_analysis_template(self, plugin, analysis):
+        try:
+            template = self._environment.get_template('{}.tex'.format(plugin))
+        except jinja2.TemplateNotFound:
+            logging.warning('Falling back on generic template for {}'.format(plugin))
+            template = self._environment.get_template(GENERIC_TEMPLATE)
+        return template.render(selected_analysis=analysis, tmp_dir=self._tmp_dir)

pdf_generator/pdf_generator.py → report.py

@@ -2,16 +2,10 @@
 
 import argparse
 import logging
-import os
-import shutil
 import sys
-from pathlib import Path
-from tempfile import TemporaryDirectory
 
 from common_helper_files import create_dir_for_file
-from common_helper_process import execute_shell_command_get_return_code
-from pdf_generator.pre_processing.rest import create_request_url, request_firmware_data
-from pdf_generator.tex_generation.template_engine import create_jinja_environment
+from pdf_generator.generator import generate_report
 
 PROGRAM_NAME = 'FACT PDF Report Generator'
 PROGRAM_VERSION = '0.1'
@@ -24,7 +18,7 @@ def setup_argparser():
     parser.add_argument('-l', '--log_file', help='path to log file')
     parser.add_argument('-L', '--log_level', help='define the log level [DEBUG,INFO,WARNING,ERROR]')
     parser.add_argument('-d', '--debug', action='store_true', default=False, help='print debug messages')
-    parser.add_argument('-u', '--uid', help='firmware analysis UID', dest="UID")
+    parser.add_argument('UID', help='firmware analysis UID')
     return parser.parse_args()
 
 
@@ -45,50 +39,7 @@ def setup_logging(debug_flag=False):
     logger.addHandler(console_log)
 
 
-def execute_pdflatex(tmp_dir):
-    current_dir = os.getcwd()
-    os.chdir(tmp_dir)
-    logging.debug('Creating pdf file')
-    _, _ = execute_shell_command_get_return_code('env buf_size=1000000 pdflatex main.tex')
-    os.chdir(current_dir)
-
-
-def _copy_fact_image(target):
-    shutil.copy(str(Path(__file__).parent.parent / 'templates' / 'fact_logo.png'), str(Path(target) / 'fact_logo.png'))
-
-
-def main(firmware_uid):
-    request_url = create_request_url(firmware_uid)
-    try:
-        firmware_analyses, firmware_meta_data = request_firmware_data(request_url)
-    except KeyError:
-        logging.warning('No firmware found with UID {}'.format(firmware_uid))
-        return None
-
-    environment = create_jinja_environment()
-
-    with TemporaryDirectory() as tmp_dir:
-        Path(tmp_dir, 'meta.tex').write_text(generate_meta_data_code(environment=environment, meta_data=firmware_meta_data))
-
-        for filename, result_code in generate_analysis_codes(environment=environment, analysis=firmware_analyses, tmp_dir=tmp_dir):
-            Path(tmp_dir, filename).write_text(result_code)
-
-        Path(tmp_dir, 'main.tex').write_text(generate_main_code(firmware_analyses, firmware_meta_data, environment))
-
-        _copy_fact_image(tmp_dir)
-
-        execute_pdflatex(tmp_dir)
-
-        pdf_filename = create_report_filename(firmware_meta_data)
-        shutil.move(str(Path(tmp_dir, 'main.pdf')), str(Path('.', pdf_filename)))
-
-    return None
-
-
 if __name__ == '__main__':
-    args = setup_argparser()
-    setup_logging(args.debug)
-
-    main(args.UID)
-
-    sys.exit(0)
+    ARGS = setup_argparser()
+    setup_logging(ARGS.debug)
+    sys.exit(generate_report(ARGS.UID))

test/acceptance/test_main.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import pytest
 from common_helper_process.fail_safe_subprocess import execute_shell_command_get_return_code
 
-SCRIPT_PATH = Path(__file__).parent.parent.parent / 'pdf_generator' / 'pdf_generator.py'
+SCRIPT_PATH = Path(__file__).parent.parent.parent / 'report.py'
 
 
 @pytest.mark.parametrize('arguments, expected_output, expected_return_code', [
@@ -11,7 +11,6 @@ SCRIPT_PATH = Path(__file__).parent.parent.parent / 'pdf_generator' / 'pdf_gener
     ('-h', 'usag', 0)
 ])
 def test_main_program(arguments, expected_output, expected_return_code):
-    command_line = str(SCRIPT_PATH) + ' ' + arguments
-    output, return_code = execute_shell_command_get_return_code(command_line)
+    output, return_code = execute_shell_command_get_return_code('{} {}'.format(SCRIPT_PATH, arguments))
     assert output[0:4] == expected_output
     assert return_code == expected_return_code

test/integration/test_code_generation.py

-from pathlib import Path
-
 import pytest
-from jinja2 import Environment, FileSystemLoader
 from pdf_generator.pre_processing.rest import request_firmware_data
-from pdf_generator.tex_generation.code_generation import generate_meta_data_code
-from pdf_generator.tex_generation.template_engine import _add_filters_to_jinja
+from pdf_generator.tex_generation.template_engine import Engine
 from test.data.test_dict import TEST_DICT
 
 # pylint: disable=redefined-outer-name
@@ -17,22 +13,8 @@ class MockResponse:
 
 
 @pytest.fixture(scope='function')
-def mock_environment():
-    env = Environment(
-        block_start_string=r'\BLOCK{',
-        block_end_string='}',
-        variable_start_string=r'\VAR{',
-        variable_end_string='}',
-        comment_start_string=r'\#{',
-        comment_end_string='}',
-        line_statement_prefix='%%',
-        line_comment_prefix='%#',
-        trim_blocks=True,
-        autoescape=False,
-        loader=FileSystemLoader(str(Path(Path(__file__).parent.parent.parent, 'pdf_generator', 'templates', 'default'))),
-    )
-    _add_filters_to_jinja(env)
-    return env
+def stub_engine():
+    return Engine()
 
 
 def test_anything_mocked(monkeypatch):
@@ -42,6 +24,6 @@ def test_anything_mocked(monkeypatch):
     assert anything
 
 
-def test_latex_code_generation(mock_environment):
-    result = generate_meta_data_code(mock_environment, TEST_DICT)
+def test_latex_code_generation(stub_engine: Engine):
+    result = stub_engine.render_meta_template(TEST_DICT)
     assert result

test/unit/test_code_generation.py

 from tempfile import TemporaryDirectory
 
-from pdf_generator.pdf_generator import main
-from pdf_generator.tex_generation.code_generation import _render_analysis_result
-from pdf_generator.tex_generation.template_engine import create_jinja_environment
+from pdf_generator.generator import generate_report
+from pdf_generator.tex_generation.template_engine import Engine
 
 TEST_DATA = {
     'analysis': {'file_hashes': {'ssdeep': 'bla', 'sha1': 'blah'}},
@@ -12,14 +11,14 @@ TEST_DATA = {
 
 def test_render_template():
     with TemporaryDirectory() as tmp_dir:
+        engine = Engine(template_folder='test', tmp_dir=tmp_dir)
         test_data = {'meta_data': '123', 'analysis': '456'}
-        jinja_env = create_jinja_environment(templates_to_use='test')
-        output = _render_analysis_result(test_data, jinja_env, 'render_test', tmp_dir)
+        output = engine.render_analysis_template(plugin='render_test', analysis=test_data)
 
     assert output == 'Test  - '
 
 
 def test_main(monkeypatch):
-    monkeypatch.setattr('pdf_generator.pdf_generator.request_firmware_data', lambda *_: (TEST_DATA['analysis'], TEST_DATA['meta_data']))
-    monkeypatch.setattr('pdf_generator.pdf_generator.shutil.move', lambda *_: None)
-    main(firmware_uid='b79ea608e2f0390744642bad472f8d9fd7e4713791857da5d5fcabf70a009e50_29626948')
+    monkeypatch.setattr('pdf_generator.generator.request_firmware_data', lambda *_: (TEST_DATA['analysis'], TEST_DATA['meta_data']))
+    monkeypatch.setattr('pdf_generator.generator.shutil.move', lambda *_: None)
+    generate_report(firmware_uid='b79ea608e2f0390744642bad472f8d9fd7e4713791857da5d5fcabf70a009e50_29626948')