| """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") |
| |
| 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)), |
| ] |
|
|
