Update main.py with reverse lookup

Allows `multi_part_header_dict` to perform a footer style lookup, this is good for files with small primary headers that score low confidence to get an aggregate score from a fixed footer (that may also be small). The combined score improves confidence.
cdgriffith · Mar 3, 2024 · a557bd5 · a557bd5
1 parent 65f1bbf
commit a557bd5
Showing 1 changed file with 27 additions and 13 deletions.
diff --git a/puremagic/main.py b/puremagic/main.py
@@ -151,20 +151,34 @@ def _identify_all(header: bytes, footer: bytes, ext=None) -> List[PureMagicWithC
     for matched in matches:
         if matched.byte_match in multi_part_header_dict:
             for magic_row in multi_part_header_dict[matched.byte_match]:
-                start = magic_row.offset
-                end = magic_row.offset + len(magic_row.byte_match)
-                if end > len(header):
-                    continue
-                if header[start:end] == magic_row.byte_match:
-                    new_matches.add(
-                        PureMagic(
-                            byte_match=header[matched.offset : end],
-                            offset=magic_row.offset,
-                            extension=magic_row.extension,
-                            mime_type=magic_row.mime_type,
-                            name=magic_row.name,
+                if '-' in str(magic_row.offset):
+                    start = magic_row.offset
+                    end = magic_row.offset + len(magic_row.byte_match)
+                    match_area = footer[start:end] if end != 0 else footer[start:]
+                    if match_area == magic_row.byte_match:
+                        new_matches.add(
+                            PureMagic(
+                                byte_match=matched.byte_match + magic_row.byte_match, 
+                                offset=magic_row.offset,
+                                extension=magic_row.extension,
+                                mime_type=magic_row.mime_type,
+                                name=magic_row.name, )
+                    )   
+                else:           
+                    start = magic_row.offset
+                    end = magic_row.offset + len(magic_row.byte_match)
+                    if end > len(header):
+                        continue
+                    if header[start:end] == magic_row.byte_match:
+                        new_matches.add(
+                            PureMagic(
+                                byte_match=header[matched.offset : end],
+                                offset=magic_row.offset,
+                                extension=magic_row.extension,
+                                mime_type=magic_row.mime_type,
+                                name=magic_row.name,
                         )
-                    )
+                    )   
 
     matches.extend(list(new_matches))
     return _confidence(matches, ext)