import hashlib, struct, sys

PAD = bytes([0x28,0xBF,0x4E,0x5E,0x4E,0x75,0x8A,0x41,0x64,0x00,0x4E,0x56,0xFF,0xFA,
             0x01,0x08,0x2E,0x2E,0x00,0xB6,0xD0,0x68,0x3E,0x80,0x2F,0x0C,0xA9,0xFE,
             0x64,0x53,0x69,0x7A])

def rc4(key, data):
    S = list(range(256)); j = 0
    for i in range(256):
        j = (j + S[i] + key[i % len(key)]) & 0xff
        S[i], S[j] = S[j], S[i]
    out = bytearray(); i = j = 0
    for b in data:
        i = (i + 1) & 0xff
        j = (j + S[i]) & 0xff
        S[i], S[j] = S[j], S[i]
        out.append(b ^ S[(S[i] + S[j]) & 0xff])
    return bytes(out)

def parse_literal_string(buf, start):
    # buf[start] == ord('(') ; returns (bytes, index_after_closing_paren)
    assert buf[start:start+1] == b'('
    i = start + 1; depth = 1; out = bytearray()
    while i < len(buf):
        c = buf[i]
        if c == 0x5c:  # backslash
            n = buf[i+1]
            simple = {0x6e:0x0a,0x72:0x0d,0x74:0x09,0x62:0x08,0x66:0x0c,
                      0x28:0x28,0x29:0x29,0x5c:0x5c}
            if n in simple:
                out.append(simple[n]); i += 2
            elif 0x30 <= n <= 0x37:  # octal, up to 3 digits
                j = i + 1; oct_digits = b''
                while j < len(buf) and len(oct_digits) < 3 and 0x30 <= buf[j] <= 0x37:
                    oct_digits += buf[j:j+1]; j += 1
                out.append(int(oct_digits, 8) & 0xff); i = j
            elif n in (0x0a,):  # line continuation
                i += 2
            elif n == 0x0d:
                i += 3 if buf[i+2:i+3] == b'\n' else 2
            else:
                out.append(n); i += 2
        elif c == 0x28:
            depth += 1; out.append(c); i += 1
        elif c == 0x29:
            depth -= 1
            if depth == 0:
                return bytes(out), i + 1
            out.append(c); i += 1
        else:
            out.append(c); i += 1
    raise ValueError("unterminated string")

def find_key_string(buf, enc_obj_region, key):
    idx = buf.index(key, *enc_obj_region)
    # skip whitespace to '('
    j = idx + len(key)
    while buf[j:j+1] in (b' ', b'\n', b'\r', b'\t'):
        j += 1
    s, _ = parse_literal_string(buf, j)
    return s

data = open('locked.pdf', 'rb').read()
password = b'hunter2'

# Locate Encrypt object (10 0 obj)
enc_start = data.index(b'10 0 obj')
enc_end = data.index(b'endobj', enc_start)
O = find_key_string(data, (enc_start, enc_end), b'/O')
U = find_key_string(data, (enc_start, enc_end), b'/U')
# P value
p_idx = data.index(b'/P', enc_start, enc_end)
p_str = data[p_idx+2:enc_end].split(b'>>')[0].strip().split()[0]
P = int(p_str)
assert len(O) == 32, len(O)

# File ID (first element of /ID array in trailer / xref stream dict)
id_idx = data.index(b'/ID[')
id1, _ = parse_literal_string(data, data.index(b'(', id_idx))
print("O len", len(O), "U len", len(U), "P", P, "ID1", id1, file=sys.stderr)

n = 128 // 8  # 16
# Algorithm 3.2 (R3)
m = hashlib.md5()
m.update((password + PAD)[:32])
m.update(O)
m.update(struct.pack('<i', P))
m.update(id1)
h = m.digest()
for _ in range(50):
    h = hashlib.md5(h[:n]).digest()
file_key = h[:n]

def obj_key(num, gen):
    ext = file_key + struct.pack('<I', num)[:3] + struct.pack('<I', gen)[:2]
    return hashlib.md5(ext).digest()[:min(n + 5, 16)]

# Verify against U (Algorithm 3.5, R3) to confirm password correctness
m2 = hashlib.md5(); m2.update(PAD); m2.update(id1)
u_chk = m2.digest()
k = file_key
for i in range(20):
    rk = bytes(b ^ i for b in k)
    u_chk = rc4(rk, u_chk)
ok = u_chk == U[:16]
print("User-password check:", "PASS" if ok else "FAIL (may be owner pw)", file=sys.stderr)

def extract_stream(buf, objnum):
    marker = ('%d 0 obj' % objnum).encode()
    s = buf.index(marker)
    dict_end = buf.index(b'>>', s)
    length = int(buf[buf.index(b'/Length', s)+7:dict_end].split()[0])
    st = buf.index(b'stream', s) + len(b'stream')
    if buf[st:st+2] == b'\r\n':
        st += 2
    elif buf[st:st+1] in (b'\n', b'\r'):
        st += 1
    return buf[st:st+length]

c3 = rc4(obj_key(3, 0), extract_stream(data, 3))
c5 = rc4(obj_key(5, 0), extract_stream(data, 5))
print("--- obj3 ---", file=sys.stderr); print(c3.decode('latin1'), file=sys.stderr)
print("--- obj5 ---", file=sys.stderr); print(c5.decode('latin1'), file=sys.stderr)

# Build a clean, unencrypted PDF with a classic xref table
objects = {
    1: b"<</Type/Pages/Kids[4 0 R 6 0 R]/Count 2/MediaBox[0 0 595 842]/Resources<</Font<</F1 2 0 R>>>>>>",
    2: b"<</Type/Font/Subtype/Type1/BaseFont/Helvetica>>",
    3: b"<</Length %d>>stream\n" % len(c3) + c3 + b"\nendstream",
    4: b"<</Type/Page/Parent 1 0 R/Contents 3 0 R/Resources<</Font<</F1 2 0 R>>>>/MediaBox[0 0 595 842]>>",
    5: b"<</Length %d>>stream\n" % len(c5) + c5 + b"\nendstream",
    6: b"<</Type/Page/Parent 1 0 R/Contents 5 0 R/Resources<</Font<</F1 2 0 R>>>>/MediaBox[0 0 595 842]>>",
    7: b"<</Type/Catalog/Pages 1 0 R>>",
}
out = bytearray(b"%PDF-1.7\n%\xe2\xe3\xcf\xd3\n")
offsets = {}
for num in sorted(objects):
    offsets[num] = len(out)
    out += b"%d 0 obj\n" % num + objects[num] + b"\nendobj\n"
xref_off = len(out)
count = max(objects) + 1
out += b"xref\n0 %d\n" % count
out += b"0000000000 65535 f \n"
for num in range(1, count):
    out += b"%010d 00000 n \n" % offsets[num]
out += b"trailer\n<</Size %d/Root 7 0 R/ID[(deft-test-file-id)(deft-test-file-id)]>>\n" % count
out += b"startxref\n%d\n%%%%EOF\n" % xref_off
open('unlocked.pdf', 'wb').write(out)
print("Wrote unlocked.pdf (%d bytes)" % len(out), file=sys.stderr)