#! /usr/bin/env python

#

# This tool is provided by dblatex (http://dblatex.sourceforge.net) and has

# the same copyright.

#

# It was initially developped to find out the fonts used and their size because

# as strange as it may seem, no obvious tool gives the font sizes used (pdffonts

# just lists the font objects of the PDF). The script can be improved to give

# more informations in a next release.

#

# To understand the PDF format, read:

#   * The reference:

#     http://www.adobe.com/content/dam/Adobe/en/devnet/acrobat/pdfs/

#                                                      pdf_reference_1-7.pdf

#

#   * A usefull introduction:

#     http://www.adobe.com/content/dam/Adobe/en/technology/pdfs/

#                                                      PDF_Day_A_Look_Inside.pdf

#

#

import os

import sys

import traceback

import zlib

import re

import logging

import tempfile

import shutil

import struct

import codecs

class ErrorHandler:

    def __init__(self):

        self._dump_stack = False

        self.rc = 0

    def dump_stack(self, dump=True):

        self._dump_stack = dump

    def failure_track(self, msg, rc=1):

        self.rc = rc

        print >>sys.stderr, (msg)

        if self._dump_stack:

            traceback.print_exc()

    def failed_exit(self, rc=1):

        self.failure_track(msg, rc)

        sys.exit(self.rc)

def pdfstring_is_list(data):

    return (data and data[0] == "[" and data[-1] == "]")

class PDFResolver:

    _resolver = None

    @classmethod

    def set_resolver(cls, resolver):

        cls._resolver = resolver

    @classmethod

    def get_resolver(cls):

        return cls._resolver

class PDFBaseObject:

    _log = logging.getLogger("pdfscan.base")

    def __init__(self):

        pass

    def debug(self, text):

        self._log.debug(text)

    def warning(self, text):

        self._log.warning(text)

    def error(self, text):

        self._log.error(text)

    def info(self, text):

        self._log.info(text)

class PDFFile(PDFBaseObject):

    """

    Main object that parses the PDF file and extract the objects needed for

    scanning.

    """

    _log = logging.getLogger("pdfscan.pdffile")

    def __init__(self, stream_manager=None):

        self._file = None

        self.filesize = 0

        self.startxref_pos = 0

        self.trailer = None

        self.xref_first = None

        self.xref_table = {}

        self.xref_objstm = {}

        self.objstm_objects = {}

        self.page_objects = []

        self.pdfobjects = PDFObjectGroup()

        self.stream_manager = stream_manager or StreamManager()

        # Create an publish the object resolver

        self.resolver = PDFObjectResolver(self)

        PDFResolver.set_resolver(self.resolver)

        # Create a global font manager

        self.fontmgr = FontManager({})

        # Detect the beginning of a PDF Object

        self.re_objstart = re.compile("(\d+) (\d+) obj(.*$)", re.DOTALL)

    def cleanup(self):

        self.stream_manager.cleanup()

    def load(self, filename):

        self.filesize = os.path.getsize(filename)

        self._file = open(filename, "rb")

        self.read_xref()

        self.build_final_xref()

    def find_startxref(self, offset_trailer=160):

        # Look for the first xref from the end

        offset, data = self.filesize, ""

        while not("startxref" in data) or offset == 0:

            offset = max(0, offset - offset_trailer)

            self._file.seek(offset)

            data = self._file.read(offset_trailer) + data

        m = re.search("\sstartxref\s+(\d+)\s+%%EOF", data, re.M)

        if not(m):

            self.error("Problem in PDF file: startxref not found")

            return 0

        self.startxref_pos = int(m.group(1))

        return self.startxref_pos

    def read_xref(self):

        startxref = self.find_startxref()

        xref = None

        while startxref:

            self._file.seek(startxref)

            line = self._file.readline()

            m = re.search("xref\s(.*)", line, re.M|re.DOTALL)

            if (m):

                found_xref = PDFXrefSection(self._file)

                found_xref.read_table(m.group(1))

            elif self.re_objstart.search(line):

                self.info("Xref section not found. Try to load XRef object")

                pdfobject, remain_line = self._parse_object(startxref)

                found_xref = PDFXrefObject(pdfobject)

            startxref = int(found_xref.trailer.get("/Prev", 0))

            if xref: xref.set_older(found_xref)

            xref = found_xref

        self.xref_first = xref

    def build_final_xref(self):

        xref = self.xref_first

        while xref:

            self.trailer = xref.trailer

            self.xref_table.update(xref.table)

            self.xref_objstm.update(xref.objstm)

            xref = xref.newer

    def get_objstm(self, objstm_id):

        return self.objstm_objects.get(objstm_id, None)

    def create_objstm(self, pdfobject):

        self.debug("Create objstm %s" % pdfobject.ident())

        pdfobject.compute()

        pdfobject.stream_decode()

        self.pdfobjects.add_object(pdfobject)

        objstm = PDFObjectStream(pdfobject)

        self.objstm_objects[objstm.ident()] = objstm

        return objstm

    def xref_resolve_object(self, ident):

        offset = self.xref_table.get(ident, 0)

        if offset != 0:

            #print "Object '%s' found at offset: %d" % (ident, offset)

            pdfobject, remain_line = self._parse_object(offset)

            return pdfobject

    def xref_resolve(self, ident):

        # Try to resolve a standard object

        pdfobject = self.xref_resolve_object(ident)

        if pdfobject:

            return pdfobject

        # Find the ObjStm infos that contains that object

        objstm_data = self.xref_objstm.get(ident, 0)

        if objstm_data == 0:

            self.warning("ObjStm id for '%s' not found in xref table" % ident)

            return None

        # If the ObjStm itself is not resolved, resolve it first

        objstm_id = "%d 0" % objstm_data[0]

        object_idx = objstm_data[1]

        objstm = self.get_objstm(objstm_id)

        if not(objstm):

            pdfobject = self.xref_resolve_object(objstm_id)

            if pdfobject: objstm = self.create_objstm(pdfobject)

        if not(objstm):

            self.error("Object '%s' cannot be resolved: ObjStm '%s' not found" \

                      % (ident, objstm_id))

            return None

        # Ok, now get the object from the ObjStm

        pdfobject = objstm.get_object(object_idx)

        return pdfobject

    def resolve_object(self, ident):

        pdfobject = self.pdfobjects.get_object(ident)

        if not(pdfobject):

            #print "Try to resolve object '%s'" % ident

            pdfobject = self.xref_resolve(ident)

            if pdfobject:

                self.pdfobjects.add_object(pdfobject)

        return pdfobject

    def get_object(self, ident):

        ident = ident.replace(" R", "").strip()

        pdfobject = self.resolver.get(ident)

        if pdfobject:

            pdfobject.link_to(self.resolver)

        return pdfobject

    def _parse_object(self, offset):

        pdfobj = None

        parsed_object = None

        remain_line = ""

        self._file.seek(offset)

        while not(parsed_object):

            line = self._file.readline()

            if not(line):

                break

            while line:

                if pdfobj:

                    fields = line.split("endobj", 1)

                    if len(fields) > 1:

                        if fields[0]:

                            pdfobj.append_string(fields[0])

                        pdfobj.compute()

                        remain_line = fields[1]

                        parsed_object = pdfobj

                    else:

                        pdfobj.append_string(line)

                    line = ""

                else:

                    m = self.re_objstart.search(line)

                    if m:

                        number, revision = m.group(1), m.group(2)

                        pdfobj = PDFObject(number, revision,

                                       stream_manager=self.stream_manager)

                        line = m.group(3)

                    else:

                        # drop the line

                        line = ""

        return (parsed_object, remain_line)

    def _expand_pages(self, page_kids):

        # Iterations to make a list of unitary pages (/Page) from a list

        # containing group of pages (/Pages). The iterations stop when all

        # The objects in the list are replaced by unit pages and not

        # intermediate page groups

        page_list = page_kids

        has_kid = len(page_list)

        while has_kid:

            newlist = []

            has_kid = 0

            for kid in page_list:

                #print kid

                kid.link_to(self.resolver)

                if kid.get_type() == "/Pages":

                    kids = kid.descriptor.get("/Kids")

                    self.debug("Expand page list: %s -> %s" % (kid, kids))

                    has_kid += len(kids)

                elif kid.get_type() == "/Page":

                    kids = [kid]

                else:

                    self.error("%s: %s" % (kid, kid.descriptor.params))

                    self.error("%s: What's wrong? '%s'" % (kid, kid.get_type()))

                    kids = []

                newlist = newlist + kids

            page_list = newlist

        return page_list

    def load_pages(self):

        root = self.trailer.get("/Root")

        catalog = self.get_object(root)

        pages = catalog.descriptor.get("/Pages")

        page_count = int(pages.descriptor.get("/Count"))

        self.info("Found %d pages" % page_count)

        pages.link_to(self.resolver)

        page_kids = pages.descriptor.get("/Kids")

        self.page_objects = self._expand_pages(page_kids)

        if len(self.page_objects) != page_count:

            self.error("Unconsistent pages found: %d vs %d" % \

                  (len(self.page_objects), page_count))

class PDFObjectResolver:

    def __init__(self, pdffile):

        self.pdffile = pdffile

    def get(self, ident, default=None):

        pdfobject = self.pdffile.resolve_object(ident)

        if not(pdfobject): pdfobject = default

        return pdfobject

class PDFObjectGroup(PDFBaseObject):

    """

    Group of the PDF Objects contained in a file. This wrapper is a dictionnary

    of the objects, and consolidates the links between the objects.

    """

    _log = logging.getLogger("pdfscan.pdffile")

    def __init__(self):

        self.pdfobjects = {}

        self.objtypes = {}

        self.unresolved = []

    def count(self):

        return len(self.pdfobjects.values())

    def types(self):

        return self.objtypes.keys()

    def add_object(self, pdfobject):

        self.pdfobjects[pdfobject.ident()] = pdfobject

        objtype = pdfobject.get_type()

        if not(objtype):

            objtype = "misc"

        lst = self.objtypes.get(objtype, [])

        lst.append(pdfobject)

        self.objtypes[objtype] = lst

        self.unresolved.append(pdfobject)

    def get_objects_by_type(self, objtype):

        return self.objtypes.get(objtype, [])

    def get_object(self, ident):

        return self.pdfobjects.get(ident, None)

    def link_objects(self):

        self.debug("%d objects to resolve" % (len(self.unresolved)))

        unresolved = []

        for pdfobj in self.unresolved:

            if pdfobj.link_to(self.pdfobjects):

                unresolved.append(pdfobj)

        self.unresolved = unresolved

    def stream_decode(self):

        for pdfobj in self.pdfobjects.values():

            pdfobj.stream_decode()

class PDFPage:

    def __init__(self, pdf, page, pagenum=0):

        self.pagenum = pagenum

        self.pdf = pdf

        contents = page.descriptor.get("/Contents")

        resources = page.descriptor.get("/Resources")

        if (isinstance(resources, PDFDescriptor)):

            rsc_descriptor = resources

        else:

            rsc_descriptor = resources.descriptor

        rsc_descriptor.link_to(pdf.resolver)

        font = rsc_descriptor.get("/Font")

        if font:

            font.link_to(pdf.resolver)

            if (isinstance(font, PDFDescriptor)):

                fontdict = font.infos()

            else:

                fontdict = font.descriptor.infos()

        else:

            fontdict = {}

        if not(isinstance(contents, list)):

            contents = [contents]

        self.page = page

        self.contents = contents

        self.fontdict = fontdict

        self.fontmgr = FontManager(fontdict, pdf.fontmgr)

        self.streams = []

        self.link_to(pdf.resolver)

        self.load_streams()

    def link_to(self, resolver):

        for content in self.contents:

            content.link_to(resolver)

    def load_streams(self):

        for content in self.contents:

            stream = PDFContentStream(content, self.fontmgr)

            self.streams.append(stream)

    def find_fonts(self):

        return self.fontmgr.get_used()

class PDFXrefSection(PDFBaseObject):

    """

    Section starting by 'xref' and followed by the 'trailer'. The xref data

    contain information about how to access to objects in the file and is

    therefore a crucial part of the object resolution.

    """

    _log = logging.getLogger("pdfscan.xref")

    _re_desc = re.compile("(<<(?:(?]|(?)>(?!>))*>>)",

                          re.MULTILINE)

    def __init__(self, fd):

        self.trailer = None

        self.table = {}

        self.objstm = {}

        self._file = fd

        self.older = None

        self.newer = None

    def set_older(self, older):

        self.older = older

        older.newer = self

    def _xref_fill_entry(self, fields, obj_id):

        offset, revision, what = fields

        if what == "n":

            ident = "%d %d" % (obj_id, int(revision))

            self.table[ident] = int(offset)

    def read_table(self, linestart=""):

        line = linestart.strip() or self._file.readline()

        subsection = line.split()

        while subsection[0] != "trailer":

            start_ref = int(subsection[0])

            object_count = int(subsection[1])

            if len(subsection) == 5:

                self._xref_fill_entry(subsection[2:], start_ref)

                start_ref += 1

                object_count -= 1

            for i in range(object_count):

                line = self._file.readline()

                self._xref_fill_entry(line.split(), start_ref+i)

            line = self._file.readline()

            subsection = line.split()

        #print len(self.table.values())

        if subsection[0] == "trailer":

            data = " ".join(subsection)

        # Ensure we have a complete dictionnary

        while not(">>" in data):

            data += self._file.readline()

        m = self._re_desc.search(data)

        if not(m):

            self.error("Problem in PDF file: cannot find valid trailer")

            return

        self.trailer = PDFDescriptor(string=m.group(1))

        self.trailer.compute()

class PDFStreamHandler:

    """

    Core abstract class in charge to handle the stream of <pdfobject>

    """

    def __init__(self, pdfobject):

        self.stream_object = pdfobject

    def ident(self):

        return self.stream_object.ident()

    def debug(self, text):

        self.stream_object.debug(text)

    def warning(self, text):

        self.stream_object.warning(text)

    def error(self, text):

        self.stream_object.error(text)

    def info(self, text):

        self.stream_object.info(text)

class PDFXrefObject(PDFStreamHandler):

    """

    A specific object that contains XRef entries in binary format. It is an

    alternative to the xref section.

    """

    def __init__(self, pdfobject):

        PDFStreamHandler.__init__(self, pdfobject)

        self.trailer = pdfobject.descriptor

        self.table = {}

        self.objstm = {}

        self.older = None

        self.newer = None

        if pdfobject.descriptor.get("/Type") != "/XRef":

            self.error("Not an XRef object. Give up")

            return

        _format = pdfobject.descriptor.get("/W")

        _format = _format.replace("[", "").replace("]", "")

        self._format = [ int(f) for f in _format.split() ]

        # An /XRef object must contains a stream

        pdfobject.stream_decode()

        self.data = pdfobject.stream_text()

        self.read_table()

    def set_older(self, older):

        self.older = older

        older.newer = self

    def _xref_fill_entry(self, fields, obj_id):

        offset, revision, what = fields

        if what == "n":

            ident = "%d %d" % (obj_id, int(revision))

            self.table[ident] = int(offset)

            self.debug("Record xref entry: '%s' @ %s" % (ident, offset))

    def _xref_fill_objstm(self, fields, obj_id):

        objstm_id, obj_index = fields

        ident = "%d %d" % (obj_id, 0)

        self.objstm[ident] = (objstm_id, obj_index)

        self.debug("Record xref entry in objstm: '%s' @ %s" % \

                   (ident, fields))

    def _int_of(self, string):

        # Convert to int from bytes string that can be of any size

        m = len(string)

        d = 0

        for i, c in enumerate(string):

            d += (1 << (8*(m - i-1))) * struct.unpack("B", c)[0]

        return d

    def read_table(self, linestart=""):

        data = self.data

        fields = 3 * [0]

        entry_size = sum(self._format)

        # TODO: use /Index

        obj_id = 0

        while data:

            first = 0

            last = 0

            for i in range(3):

                last += self._format[i]

                fields[i] = self._int_of(data[first:last])

                first = last

            data = data[entry_size:]

            if fields[0] == 1:

                self._xref_fill_entry(fields[1:3] + ["n"], obj_id)

            elif fields[0] == 2:

                self._xref_fill_objstm(fields[1:3], obj_id)

            obj_id += 1

class PDFObjectStream(PDFStreamHandler):

    """

    A PDF Object Stream contains in its stream some compressed PDF objects.

    This class works on a PDF object stream to build the containded PDF objects.

    """

    def __init__(self, pdfobject):

        PDFStreamHandler.__init__(self, pdfobject)

        self._pdfobjects = []

    def pdfobjects(self):

        if not(self._pdfobjects):

            self.compute()

        return self._pdfobjects

    def _getinfo(self, what):

        return self.stream_object.descriptor.get(what)

    def get_object(self, idx):

        if not(self._pdfobjects):

            self.compute()

        if idx < 0 or idx >= len(self._pdfobjects):

            return None

        return self._pdfobjects[idx]

    def parse_object_list(self, data):

        values = data.split()

        objlist = []

        for i in range(0, len(values), 2):

            # The pair is ('object number', byte_offset)

            objlist.append((values[i], int(values[i+1])))

        self.objlist = objlist

        return objlist

    def compute(self):

        _type = self._getinfo("/Type")

        if  _type != "/ObjStm":

            self.error("Cannot read object stream: Invalid type '%s'" % _type)

            return

        nb_objects = int(self._getinfo("/N"))

        objlist_b = int(self._getinfo("/First"))

        stream = self.stream_object.stream_cache

        objlist = self.parse_object_list(stream.read(objlist_b))

        if len(objlist) != nb_objects:

            self.warning("Error in parsing the Stream Object: found %d"\

                         "objects instead of %d" % (len(objlist), nb_object))

        # List Terminator

        objlist.append(("",-1))

        bytes_read = 0

        for i in range(len(objlist)-1):

            # In ObjectStream, a PDF object revision is always '0'

            number, revision = objlist[i][0], "0"

            # The size of the object data is given by the position of the next

            objsize = objlist[i+1][1] - bytes_read

            if objsize >= 0:

                data = stream.read(objsize)

            else:

                data = stream.read()

            bytes_read += len(data)

            self.debug("Object[%d] in stream: '%s' has %d bytes" % \

                       (i, number, objsize))

            # Build the PDF Object from stream data

            pdfobj = PDFObject(number, revision)

            pdfobj.append_string(data)

            pdfobj.compute()

            self._pdfobjects.append(pdfobj)

        stream.close()

class PDFObject:

    """

    A PDF Object contains the data between the 'obj ... 'endobj' tags.

    It has a unique identifier given by the (number,revision) pair.

    The data contained by a PDF object can be dictionnaries (descriptors),

    stream contents and other stuff.

    """

    # Extract a dictionnary '<<...>>' leaf (does not contain another dict)

    _re_desc = re.compile("(<<(?:(?]|(?)>(?!>))*>>)",

                          re.MULTILINE)

    def __init__(self, number, revision, stream_manager=None):

        self.string = ""

        self.number = number

        self.revision = revision

        self.descriptors = []

        self.descriptor = None

        self.data = ""

        self.stream = None

        self.outfile = ""

        self.stream_manager = stream_manager or StreamManager()

        self._log = logging.getLogger("pdfscan.pdfobject")

        self.debug("New Object")

        self.re_desc = self._re_desc

    def debug(self, text):

        self._log.debug(self.logstr(text))

    def warning(self, text):

        self._log.warning(self.logstr(text))

    def error(self, text):

        self._log.error(self.logstr(text))

    def info(self, text):

        self._log.info(self.logstr(text))

    def ident(self):

        return "%s %s" % (self.number, self.revision)

    def __repr__(self):

        return "(%s R)" % self.ident()

    def __int__(self):

        return int(self.data)

    def logstr(self, text):

        return "Object [%s %s]: %s" % (self.number,self.revision,text)

    def append_string(self, string):

        self.string = self.string + string

    def compute(self):

        string = self.string

        s = re.split("stream\s", string, re.MULTILINE)

        if len(s) > 1:

            self.debug("Contains stream")

            self.stream = s[1].strip()

        string = s[0]

        # Iterate to build all the nested dictionnaries/descriptors,

        # from the deepest to the main one

        self.descriptors = []

        while True:

            descs = self.re_desc.findall(string)

            if not(descs):

                break

            for desc_str in descs:

                desc = PDFDescriptor(string=desc_str)

                string = string.replace(desc_str,

                            "{descriptor(%d)}" % len(self.descriptors))

                self.descriptors.append(desc)

        self.debug("Found %d descriptors" % len(self.descriptors))

        for descobj in self.descriptors:

            descobj.compute(descriptors=self.descriptors)

        if self.descriptors:

            self.descriptor = self.descriptors[-1]

        else:

            self.descriptor = PDFDescriptor()

        self.data = re.sub("{descriptor\(\d+\)}", "",

                           string, flags=re.MULTILINE).strip()

        self.debug("Data: '%s'" % self.data)

    def stream_decode(self):

        if not(self.stream):

            return

        self.debug("Try to decode stream...")

        # Consolidate stream buffer from the /Length information

        stream_size = int(self.descriptor.get("/Length"))

        self.stream = self.stream[0:stream_size]

        # Put the stream in a cache

        self.stream_cache = self.stream_manager.cache(number=self.number,

                                                      revision=self.revision)

        method = self.descriptor.get("/Filter")

        if method == "/FlateDecode":

            method = "zlib"

        elif method == "/DCTDecode":

            # This is JPEG. Just dump it

            self.warning("this is a JPEG stream")

            method = ""

        elif method != "":

            self.error("don't know how to decode stream with filter '%s'" \

                     % method)

            return

        self.stream_cache.write(self.stream, compress_type=method)

    def stream_text(self):

        if not(self.stream):

            return ""

        data = self.stream_cache.read()

        self.stream_cache.close()

        return data

    def get_type(self):

        _type = self.descriptor.get("/Type")

        if _type:

            return _type

        if self.stream:

            return "stream"

        if pdfstring_is_list(self.data):

            return "list"

        if self.descriptor.is_name_tree_node():

            return "name tree"

    def link_to(self, pdfobjects):

        self.debug("Link objects")

        for desc in self.descriptors:

            desc.link_to(pdfobjects)

        if pdfstring_is_list(self.data):

            pass

class PDFDescriptor:

    """

    Contains the data between the << ... >> brackets in PDF objects. It is

    a dictionnary that can contain other descriptors/dictionnaries.

    """

    # Unique identifier for these objects

    _id = 0

    # Detect the dictionnary fields covering these cases:

    # <<

    #  /Type /Page                    : the value is another keyword

    #  /Contents 5 0 R                : the value is a string up next keyword

    #  /Resources 4 0 R                   

    #  /MediaBox [0 0 595.276 841.89] : the value is an array

    #  /Parent 12 0 R

    # >>

    _re_dict = re.compile("/\w+\s*/[^/\s]+|/\w+\s*\[[^\]]*\]|/\w+\s*[^/]+")

    # Extract a dictionnary keyword

    _re_key = re.compile("(/[^ \({/\[<]*)")

    # Extract the substituted descriptors

    _re_descobj = re.compile("{descriptor\((\d+)\)}")

    # Find the PDF object references

    _re_objref = re.compile("(\d+ \d+ R)")

    def __init__(self, string=""):

        self._ident = self._get_ident()

        self.string = string

        self.params = {}

        self._log = logging.getLogger("pdfscan.descriptor")

        self.re_dict = self._re_dict

        self.re_key = self._re_key

        self.re_descobj = self._re_descobj

        self.re_objref = self._re_objref

    def _get_ident(self):

        _id = PDFDescriptor._id

        PDFDescriptor._id += 1

        return _id

    def ident(self):

        return self._ident

    def debug(self, text):

        self._log.debug("Descriptor [%d]: %s" % (self._ident, text))

    def error(self, text):

        self._log.error("Descriptor [%d]: %s" % (self._ident, text))

    def info(self, text):

        self._log.info("Descriptor [%d]: %s" % (self._ident, text))

    def warning(self, text):

        self._log.warning("Descriptor [%d]: %s" % (self._ident, text))

    def __repr__(self):

        return "desc[%d]" % self._ident

    def normalize_fields(self, string):

        string = string.replace(">>", "")

        string = string.replace("<<", "")

        string = string.replace("\n", " ")

        fields = self.re_dict.findall(string)

        fields = [ f.strip() for f in fields if (f and f.strip()) ]

        return fields

    def compute(self, descriptors=None):

        lines = self.normalize_fields(self.string)

        for line in lines:

            m = self.re_key.match(line)

            if not(m):

                continue

            param = m.group(1)

            value = line.replace(param, "").strip()

            m = self.re_descobj.match(value)

            if m and descriptors:

                value = descriptors[int(m.group(1))]

            self.params[param] = value

        self.debug(self.params)

    def get(self, param, default=""):

        return self.params.get(param, default)

    def values(self):

        return self.params.values()

    def keys(self):

        return self.params.keys()