File size: 2,584 Bytes
f954ee9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | """Byte-level structural features used by pdf-damage-classifier.
The same function is used for training and inference, so the ONNX model's
input vector is defined by FEATURE_NAMES (in this exact order).
Only the PDF's raw bytes are needed; no PDF parser is involved, so it works
on files that are too broken to open.
"""
import math
import re
FEATURE_NAMES = [
"has_pdf_header",
"eof_in_last_1k",
"ends_with_eof",
"startxref_present",
"startxref_numeric",
"startxref_points_to_xref",
"xref_valid_fraction",
"has_trailer",
"root_object_exists",
"obj_endobj_gap",
"stream_endstream_gap",
"log10_size",
]
_XREF_ENTRY = re.compile(rb"(\d{10}) (\d{5}) n")
_STARTXREF = re.compile(rb"startxref\s+(\S+)")
_ROOT = re.compile(rb"/Root\s+(\d+)\s+(\d+)\s+R")
def extract_features(data: bytes):
n = len(data)
tail = data[-1024:]
has_header = float(data[:5] == b"%PDF-")
eof_in_tail = float(b"%%EOF" in tail)
ends_eof = float(data.rstrip()[-5:] == b"%%EOF")
m = None
for m in _STARTXREF.finditer(tail):
pass
sx_present = float(m is not None)
sx_numeric = 0.0
sx_points = 0.0
if m is not None:
tok = m.group(1)
if tok.isdigit():
sx_numeric = 1.0
off = int(tok)
if off < n and data[off : off + 4] == b"xref":
sx_points = 1.0
entries = _XREF_ENTRY.findall(data)
if entries:
ok = 0
for off_s, _gen in entries:
off = int(off_s)
if off < n and re.match(rb"\d+ \d+ obj", data[off : off + 24]):
ok += 1
xref_frac = ok / len(entries)
else:
xref_frac = 0.0
has_trailer = float(b"trailer" in data)
rm = _ROOT.search(data)
root_exists = 0.0
if rm is not None:
root_exists = float(re.search(rb"(?:^|\s)" + rm.group(1) + b" " + rm.group(2) + b" obj", data) is not None)
objs = len(re.findall(rb"\d+ \d+ obj\b", data))
endobjs = data.count(b"endobj")
# Count stream starts, ignoring the "stream\n" inside "endstream\n".
stream_starts = len(re.findall(rb"(?<!end)stream\r?\n", data))
endstreams = data.count(b"endstream")
obj_gap = (objs - endobjs) / max(objs, 1)
stream_gap = (stream_starts - endstreams) / max(stream_starts, 1)
return [
has_header,
eof_in_tail,
ends_eof,
sx_present,
sx_numeric,
sx_points,
xref_frac,
has_trailer,
root_exists,
obj_gap,
stream_gap,
math.log10(max(n, 1)),
]
|