Merge pull request #36 from petrs/cc_fault_tolerant_extraction

Adds fault tolerant data extraction

3 files changed, 153 insertions, 120 deletions

diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 6f1c35b9..a3b72a37 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -1096,35 +1096,52 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
 
     @staticmethod
     def extract_st_pdf_metadata(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
-        cert.pdf_data.st_metadata = helpers.extract_pdf_metadata(cert.state.st_pdf_path)[1]
+        response, cert.pdf_data.st_metadata = helpers.extract_pdf_metadata(cert.state.st_pdf_path)
+        if response != constants.RETURNCODE_OK:
+            cert.state.st_extract_ok = False
         return cert
 
     @staticmethod
     def extract_report_pdf_metadata(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
-        cert.pdf_data.report_metadata = helpers.extract_pdf_metadata(cert.state.report_pdf_path)[1]
+        response, cert.pdf_data.report_metadata = helpers.extract_pdf_metadata(cert.state.report_pdf_path)
+        if response != constants.RETURNCODE_OK:
+            cert.state.report_extract_ok = False
         return cert
 
     @staticmethod
     def extract_st_pdf_frontpage(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
         cert.pdf_data.st_frontpage = dict()
-        cert.pdf_data.st_frontpage['bsi'] = helpers.search_only_headers_bsi(cert.state.st_txt_path)
-        cert.pdf_data.st_frontpage['anssi'] = helpers.search_only_headers_anssi(cert.state.st_txt_path)
+
+        response_bsi, cert.pdf_data.st_frontpage['bsi'] = helpers.search_only_headers_bsi(cert.state.st_txt_path)
+        response_anssi, cert.pdf_data.st_frontpage['anssi'] = helpers.search_only_headers_anssi(cert.state.st_txt_path)
+
+        if response_anssi != constants.RETURNCODE_OK or response_bsi != constants.RETURNCODE_OK:
+            cert.state.st_extract_ok = False
+
         return cert
 
     @staticmethod
     def extract_report_pdf_frontpage(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
         cert.pdf_data.report_frontpage = dict()
-        cert.pdf_data.report_frontpage['bsi'] = helpers.search_only_headers_bsi(cert.state.report_txt_path)
-        cert.pdf_data.report_frontpage['anssi'] = helpers.search_only_headers_anssi(cert.state.report_txt_path)
+        response_bsi, cert.pdf_data.report_frontpage['bsi'] = helpers.search_only_headers_bsi(cert.state.report_txt_path)
+        response_anssi, cert.pdf_data.report_frontpage['anssi'] = helpers.search_only_headers_anssi(cert.state.report_txt_path)
+
+        if response_anssi != constants.RETURNCODE_OK or response_bsi != constants.RETURNCODE_OK:
+            cert.state.report_extract_ok = False
+
         return cert
 
     @staticmethod
     def extract_report_pdf_keywords(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
-        cert.pdf_data.report_keywords = helpers.extract_keywords(cert.state.report_txt_path)
+        response, cert.pdf_data.report_keywords = helpers.extract_keywords(cert.state.report_txt_path)
+        if response != constants.RETURNCODE_OK:
+            cert.state.report_extract_ok = False
         return cert
 
     @staticmethod
     def extract_st_pdf_keywords(cert: 'CommonCriteriaCert') -> 'CommonCriteriaCert':
-        cert.pdf_data.st_keywords = helpers.extract_keywords(cert.state.st_txt_path)
+        response, cert.pdf_data.st_keywords = helpers.extract_keywords(cert.state.st_txt_path)
+        if response != constants.RETURNCODE_OK:
+            cert.state.st_extract_ok = False
         return cert
 
diff --git a/sec_certs/constants.py b/sec_certs/constants.py
index 1923937c..7f63d7e9 100644
--- a/sec_certs/constants.py
+++ b/sec_certs/constants.py
@@ -3,6 +3,7 @@ from enum import Enum
 N_THREADS = 8
 RESPONSE_OK = 200
 RETURNCODE_OK = 0
+RETURNCODE_NOK = -1
 REQUEST_TIMEOUT = 5
 
 MIN_CORRECT_CERT_SIZE = 5000
diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
index bbb1be42..20d4fde3 100644
--- a/sec_certs/helpers.py
+++ b/sec_certs/helpers.py
@@ -190,8 +190,9 @@ def extract_pdf_metadata(filepath: Path):
 
     except Exception as e:
         logger.error(f'Failed to read metadata of {filepath}, error: {e}')
+        return constants.RETURNCODE_NOK, None
 
-    return filepath, metadata
+    return constants.RETURNCODE_OK, metadata
 
 
 # TODO: Please, refactor me. I reallyyyyyyyyyyyyy need it!!!!!!
@@ -292,74 +293,80 @@ def search_only_headers_anssi(filepath: Path):
          'Référence du rapport de certification(.+)Nom du produit(.+)Référence/version du produit(.+)Critères d\'évaluation et version(.+)Niveau d\'évaluation(.+)Développeurs(.+)Centre d\'évaluation(.+)Accords de reconnaissance applicables'),
     ]
 
+
     # statistics about rules success rate
     num_rules_hits = {}
     for rule in rules_certificate_preface:
         num_rules_hits[rule[1]] = 0
 
     items_found = {}
-    whole_text, whole_text_with_newlines, was_unicode_decode_error = extract_certificates.load_cert_file(filepath)
-
-    # for ANSII and DCSSI certificates, front page starts only on third page after 2 newpage signs
-    pos = whole_text.find('')
-    if pos != -1:
-        pos = whole_text.find('', pos)
-        if pos != -1:
-            whole_text = whole_text[pos:]
 
-    no_match_yet = True
-    other_rule_already_match = False
-    rule_index = -1
-    for rule in rules_certificate_preface:
-        rule_index += 1
-        rule_and_sep = rule[1] + REGEXEC_SEP
+    try:
+        whole_text, whole_text_with_newlines, was_unicode_decode_error = extract_certificates.load_cert_file(filepath)
 
-        for m in re.finditer(rule_and_sep, whole_text):
-            if no_match_yet:
-                items_found[constants.TAG_HEADER_MATCH_RULES] = []
-                no_match_yet = False
+        # for ANSII and DCSSI certificates, front page starts only on third page after 2 newpage signs
+        pos = whole_text.find('')
+        if pos != -1:
+            pos = whole_text.find('', pos)
+            if pos != -1:
+                whole_text = whole_text[pos:]
 
-            # insert rule if at least one match for it was found
-            if rule not in items_found[constants.TAG_HEADER_MATCH_RULES]:
-                items_found[constants.TAG_HEADER_MATCH_RULES].append(rule[1])
+        no_match_yet = True
+        other_rule_already_match = False
+        rule_index = -1
+        for rule in rules_certificate_preface:
+            rule_index += 1
+            rule_and_sep = rule[1] + REGEXEC_SEP
 
-            if not other_rule_already_match:
-                other_rule_already_match = True
-            else:
-                logger.warning(f'WARNING: multiple rules are matching same certification document: {filepath}')
+            for m in re.finditer(rule_and_sep, whole_text):
+                if no_match_yet:
+                    items_found[constants.TAG_HEADER_MATCH_RULES] = []
+                    no_match_yet = False
 
-            num_rules_hits[rule[1]] += 1  # add hit to this rule
-            match_groups = m.groups()
-            index_next_item = 0
-            items_found[constants.TAG_CERT_ID] = extract_certificates.normalize_match_string(match_groups[index_next_item])
-            index_next_item += 1
+                # insert rule if at least one match for it was found
+                if rule not in items_found[constants.TAG_HEADER_MATCH_RULES]:
+                    items_found[constants.TAG_HEADER_MATCH_RULES].append(rule[1])
 
-            items_found[constants.TAG_CERT_ITEM] = extract_certificates.normalize_match_string(match_groups[index_next_item])
-            index_next_item += 1
+                if not other_rule_already_match:
+                    other_rule_already_match = True
+                else:
+                    logger.warning(f'WARNING: multiple rules are matching same certification document: {filepath}')
 
-            if rule[0] == HEADER_TYPE.HEADER_MISSING_CERT_ITEM_VERSION:
-                items_found[constants.TAG_CERT_ITEM_VERSION] = ''
-            else:
-                items_found[constants.TAG_CERT_ITEM_VERSION] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                num_rules_hits[rule[1]] += 1  # add hit to this rule
+                match_groups = m.groups()
+                index_next_item = 0
+                items_found[constants.TAG_CERT_ID] = extract_certificates.normalize_match_string(match_groups[index_next_item])
                 index_next_item += 1
 
-            if rule[0] == HEADER_TYPE.HEADER_MISSING_PROTECTION_PROFILES:
-                items_found[constants.TAG_REFERENCED_PROTECTION_PROFILES] = ''
-            else:
-                items_found[constants.TAG_REFERENCED_PROTECTION_PROFILES] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                items_found[constants.TAG_CERT_ITEM] = extract_certificates.normalize_match_string(match_groups[index_next_item])
                 index_next_item += 1
 
-            items_found[constants.TAG_CC_VERSION] = extract_certificates.normalize_match_string(match_groups[index_next_item])
-            index_next_item += 1
+                if rule[0] == HEADER_TYPE.HEADER_MISSING_CERT_ITEM_VERSION:
+                    items_found[constants.TAG_CERT_ITEM_VERSION] = ''
+                else:
+                    items_found[constants.TAG_CERT_ITEM_VERSION] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                    index_next_item += 1
+
+                if rule[0] == HEADER_TYPE.HEADER_MISSING_PROTECTION_PROFILES:
+                    items_found[constants.TAG_REFERENCED_PROTECTION_PROFILES] = ''
+                else:
+                    items_found[constants.TAG_REFERENCED_PROTECTION_PROFILES] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                    index_next_item += 1
+
+                items_found[constants.TAG_CC_VERSION] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                index_next_item += 1
 
-            items_found[constants.TAG_CC_SECURITY_LEVEL] = extract_certificates.normalize_match_string(match_groups[index_next_item])
-            index_next_item += 1
+                items_found[constants.TAG_CC_SECURITY_LEVEL] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                index_next_item += 1
 
-            items_found[constants.TAG_DEVELOPER] = extract_certificates.normalize_match_string(match_groups[index_next_item])
-            index_next_item += 1
+                items_found[constants.TAG_DEVELOPER] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                index_next_item += 1
 
-            items_found[constants.TAG_CERT_LAB] = extract_certificates.normalize_match_string(match_groups[index_next_item])
-            index_next_item += 1
+                items_found[constants.TAG_CERT_LAB] = extract_certificates.normalize_match_string(match_groups[index_next_item])
+                index_next_item += 1
+    except Exception as e:
+        logger.error(f'Failed to parse ANSSI frontpage headers from {filepath}; {e}')
+        return constants.RETURNCODE_NOK, None
 
     # if True:
     #     print('# hits for rule')
@@ -371,7 +378,7 @@ def search_only_headers_anssi(filepath: Path):
     #         if rule[1] > 0:
     #             used_rules.append(rule[0])
 
-    return items_found
+    return constants.RETURNCODE_OK, items_found
 
 # TODO: Please refactor me. I need it so badlyyyyyy!!!
 def search_only_headers_bsi(filepath: Path):
@@ -384,81 +391,89 @@ def search_only_headers_bsi(filepath: Path):
 
     items_found = {}
     no_match_yet = True
-    #
-    # Process front page with info: cert_id, certified_item and developer
-    #
-    whole_text, whole_text_with_newlines, was_unicode_decode_error = extract_certificates.load_cert_file(filepath, NUM_LINES_TO_INVESTIGATE, LINE_SEPARATOR_STRICT)
 
-    for rule in rules_certificate_preface:
-        rule_and_sep = rule + REGEXEC_SEP
+    try:
+        # Process front page with info: cert_id, certified_item and developer
+        whole_text, whole_text_with_newlines, was_unicode_decode_error = extract_certificates.load_cert_file(filepath, NUM_LINES_TO_INVESTIGATE, LINE_SEPARATOR_STRICT)
+
+        for rule in rules_certificate_preface:
+            rule_and_sep = rule + REGEXEC_SEP
 
-        for m in re.finditer(rule_and_sep, whole_text):
-            if no_match_yet:
-                items_found[constants.TAG_HEADER_MATCH_RULES] = []
-                no_match_yet = False
+            for m in re.finditer(rule_and_sep, whole_text):
+                if no_match_yet:
+                    items_found[constants.TAG_HEADER_MATCH_RULES] = []
+                    no_match_yet = False
 
-            # insert rule if at least one match for it was found
-            if rule not in items_found[constants.TAG_HEADER_MATCH_RULES]:
-                items_found[constants.TAG_HEADER_MATCH_RULES].append(rule)
+                # insert rule if at least one match for it was found
+                if rule not in items_found[constants.TAG_HEADER_MATCH_RULES]:
+                    items_found[constants.TAG_HEADER_MATCH_RULES].append(rule)
 
-            match_groups = m.groups()
-            cert_id = match_groups[0]
-            certified_item = match_groups[1]
-            developer = match_groups[2]
+                match_groups = m.groups()
+                cert_id = match_groups[0]
+                certified_item = match_groups[1]
+                developer = match_groups[2]
 
-            FROM_KEYWORD_LIST = [' from ', ' der ']
-            for from_keyword in FROM_KEYWORD_LIST:
-                from_keyword_len = len(from_keyword)
-                if certified_item.find(from_keyword) != -1:
-                    logger.warning(f'string {from_keyword} detected in certified item - shall not be here, fixing...')
-                    certified_item_first = certified_item[:certified_item.find(from_keyword)]
-                    developer = certified_item[certified_item.find(from_keyword) + from_keyword_len:]
-                    certified_item = certified_item_first
-                    continue
+                FROM_KEYWORD_LIST = [' from ', ' der ']
+                for from_keyword in FROM_KEYWORD_LIST:
+                    from_keyword_len = len(from_keyword)
+                    if certified_item.find(from_keyword) != -1:
+                        logger.warning(f'string {from_keyword} detected in certified item - shall not be here, fixing...')
+                        certified_item_first = certified_item[:certified_item.find(from_keyword)]
+                        developer = certified_item[certified_item.find(from_keyword) + from_keyword_len:]
+                        certified_item = certified_item_first
+                        continue
 
-            end_pos = developer.find('\f-')
-            if end_pos == -1:
-                end_pos = developer.find('\fBSI')
-            if end_pos == -1:
-                end_pos = developer.find('Bundesamt')
-            if end_pos != -1:
-                developer = developer[:end_pos]
+                end_pos = developer.find('\f-')
+                if end_pos == -1:
+                    end_pos = developer.find('\fBSI')
+                if end_pos == -1:
+                    end_pos = developer.find('Bundesamt')
+                if end_pos != -1:
+                    developer = developer[:end_pos]
 
-            items_found[constants.TAG_CERT_ID] = extract_certificates.normalize_match_string(cert_id)
-            items_found[constants.TAG_CERT_ITEM] = extract_certificates.normalize_match_string(certified_item)
-            items_found[constants.TAG_DEVELOPER] = extract_certificates.normalize_match_string(developer)
-            items_found[constants.TAG_CERT_LAB] = 'BSI'
+                items_found[constants.TAG_CERT_ID] = extract_certificates.normalize_match_string(cert_id)
+                items_found[constants.TAG_CERT_ITEM] = extract_certificates.normalize_match_string(certified_item)
+                items_found[constants.TAG_DEVELOPER] = extract_certificates.normalize_match_string(developer)
+                items_found[constants.TAG_CERT_LAB] = 'BSI'
 
-    # Process page with more detailed certificate info
-    # PP Conformance, Functionality, Assurance
-    rules_certificate_third = ['PP Conformance: (.+)Functionality: (.+)Assurance: (.+)The IT Product identified']
+        # Process page with more detailed certificate info
+        # PP Conformance, Functionality, Assurance
+        rules_certificate_third = ['PP Conformance: (.+)Functionality: (.+)Assurance: (.+)The IT Product identified']
 
-    whole_text, whole_text_with_newlines, was_unicode_decode_error = extract_certificates.load_cert_file(filepath)
+        whole_text, whole_text_with_newlines, was_unicode_decode_error = extract_certificates.load_cert_file(filepath)
 
-    for rule in rules_certificate_third:
-        rule_and_sep = rule + REGEXEC_SEP
+        for rule in rules_certificate_third:
+            rule_and_sep = rule + REGEXEC_SEP
 
-        for m in re.finditer(rule_and_sep, whole_text):
-            # check if previous rules had at least one match
-            if not constants.TAG_CERT_ID in items_found.keys():
-                logger.error('ERROR: front page not found for file: {}'.format(filepath))
+            for m in re.finditer(rule_and_sep, whole_text):
+                # check if previous rules had at least one match
+                if not constants.TAG_CERT_ID in items_found.keys():
+                    logger.error('ERROR: front page not found for file: {}'.format(filepath))
 
-            match_groups = m.groups()
-            ref_protection_profiles = match_groups[0]
-            cc_version = match_groups[1]
-            cc_security_level = match_groups[2]
+                match_groups = m.groups()
+                ref_protection_profiles = match_groups[0]
+                cc_version = match_groups[1]
+                cc_security_level = match_groups[2]
 
-            items_found[constants.TAG_REFERENCED_PROTECTION_PROFILES] = extract_certificates.normalize_match_string(ref_protection_profiles)
-            items_found[constants.TAG_CC_VERSION] = extract_certificates.normalize_match_string(cc_version)
-            items_found[constants.TAG_CC_SECURITY_LEVEL] = extract_certificates.normalize_match_string(cc_security_level)
+                items_found[constants.TAG_REFERENCED_PROTECTION_PROFILES] = extract_certificates.normalize_match_string(ref_protection_profiles)
+                items_found[constants.TAG_CC_VERSION] = extract_certificates.normalize_match_string(cc_version)
+                items_found[constants.TAG_CC_SECURITY_LEVEL] = extract_certificates.normalize_match_string(cc_security_level)
 
-    # print('\n*** Certificates without detected preface:')
-    # for file_name in files_without_match:
-    #     print('No hits for {}'.format(file_name))
-    # print('Total no hits files: {}'.format(len(files_without_match)))
-    # print('\n**********************************')
+        # print('\n*** Certificates without detected preface:')
+        # for file_name in files_without_match:
+        #     print('No hits for {}'.format(file_name))
+        # print('Total no hits files: {}'.format(len(files_without_match)))
+        # print('\n**********************************')
+    except Exception as e:
+        logger.error(f'Failed to parse BSI headers from frontpage: {filepath}; {e}')
+        return constants.RETURNCODE_NOK, None
 
-    return items_found
+    return constants.RETURNCODE_OK, items_found
 
-def extract_keywords(filepath: Path) -> Dict[str, str]:
-    return extract_certificates.parse_cert_file(filepath, cc_search_rules, -1, extract_certificates.LINE_SEPARATOR)[0]
\ No newline at end of file
+def extract_keywords(filepath: Path) -> Tuple[int, Optional[Dict[str, str]]]:
+    try:
+        result = extract_certificates.parse_cert_file(filepath, cc_search_rules, -1, extract_certificates.LINE_SEPARATOR)[0]
+    except Exception as e:
+        logger.error(f'Failed to parse keywords from: {filepath}; {e}')
+        return constants.RETURNCODE_NOK, None
+    return constants.RETURNCODE_OK, result
\ No newline at end of file