home *** CD-ROM | disk | FTP | other *** search
/ Maximum CD 2010 November / maximum-cd-2010-11.iso / DiscContents / calibre-0.7.13.msi / file_965 (.txt) < prev    next >
Encoding:
Python Compiled Bytecode  |  2010-08-06  |  14.2 KB  |  432 lines

  1. # Source Generated with Decompyle++
  2. # File: in.pyc (Python 2.6)
  3.  
  4. from __future__ import with_statement
  5. __license__ = 'GPL v3'
  6. __copyright__ = '2009, Kovid Goyal kovid@kovidgoyal.net and Marshall T. Vandegrift <llasram@gmail.com>'
  7. __docformat__ = 'restructuredtext en'
  8. from struct import pack, unpack
  9. from cStringIO import StringIO
  10. from calibre.ebooks.mobi import MobiError
  11. from calibre.ebooks.mobi.writer import rescale_image, MAX_THUMB_DIMEN
  12. from calibre.ebooks.mobi.langcodes import iana2mobi
  13. from calibre.utils.date import now as nowf
  14.  
  15. class StreamSlicer(object):
  16.     
  17.     def __init__(self, stream, start = 0, stop = None):
  18.         self._stream = stream
  19.         self.start = start
  20.         if stop is None:
  21.             stream.seek(0, 2)
  22.             stop = stream.tell()
  23.         
  24.         self.stop = stop
  25.         self._len = stop - start
  26.  
  27.     
  28.     def __len__(self):
  29.         return self._len
  30.  
  31.     
  32.     def __getitem__(self, key):
  33.         stream = self._stream
  34.         base = self.start
  35.         if isinstance(key, (int, long)):
  36.             stream.seek(base + key)
  37.             return stream.read(1)
  38.         if isinstance(key, slice):
  39.             (start, stop, stride) = key.indices(self._len)
  40.             if stride < 0:
  41.                 start = stop
  42.                 stop = start
  43.             
  44.             size = stop - start
  45.             if size <= 0:
  46.                 return ''
  47.             stream.seek(base + start)
  48.             data = stream.read(size)
  49.             if stride != 1:
  50.                 data = data[::stride]
  51.             
  52.             return data
  53.         raise TypeError('stream indices must be integers')
  54.  
  55.     
  56.     def __setitem__(self, key, value):
  57.         stream = self._stream
  58.         base = self.start
  59.         if isinstance(key, (int, long)):
  60.             if len(value) != 1:
  61.                 raise ValueError('key and value lengths must match')
  62.             len(value) != 1
  63.             stream.seek(base + key)
  64.             return stream.write(value)
  65.         if isinstance(key, slice):
  66.             (start, stop, stride) = key.indices(self._len)
  67.             if stride < 0:
  68.                 start = stop
  69.                 stop = start
  70.             
  71.             size = stop - start
  72.             if stride != 1:
  73.                 value = value[::stride]
  74.             
  75.             if len(value) != size:
  76.                 raise ValueError('key and value lengths must match')
  77.             len(value) != size
  78.             stream.seek(base + start)
  79.             return stream.write(value)
  80.         raise TypeError('stream indices must be integers')
  81.  
  82.     
  83.     def update(self, data_blocks):
  84.         stream = self._stream
  85.         base = self.start
  86.         stream.seek(base)
  87.         self._stream.truncate(base)
  88.         for block in data_blocks:
  89.             stream.write(block)
  90.         
  91.  
  92.     
  93.     def truncate(self, value):
  94.         self._stream.truncate(value)
  95.  
  96.  
  97.  
  98. class MetadataUpdater(object):
  99.     DRM_KEY_SIZE = 48
  100.     
  101.     def __init__(self, stream):
  102.         self.stream = stream
  103.         data = self.data = StreamSlicer(stream)
  104.         self.type = data[60:68]
  105.         if self.type != 'BOOKMOBI':
  106.             return None
  107.         (self.nrecs,) = unpack('>H', data[76:78])
  108.         (mobi_header_length,) = unpack('>I', record0[20:24])
  109.         if not mobi_header_length:
  110.             raise MobiError("Non-standard file format.  Try 'Convert E-Books' with MOBI as Input and Output formats.")
  111.         mobi_header_length
  112.         (self.encryption_type,) = unpack('>H', record0[12:14])
  113.         (codepage,) = unpack('>I', record0[28:32])
  114.         self.codec = record0 = self.record0 = self.record(0) if codepage == 65001 else 'cp1252'
  115.         (image_base,) = unpack('>I', record0[108:112])
  116.         self.cover_record = None
  117.         self.thumbnail_record = None
  118.         self.timestamp = None
  119.         self.pdbrecords = self.get_pdbrecords()
  120.         self.drm_block = None
  121.         if self.encryption_type != 0:
  122.             if self.have_exth:
  123.                 self.drm_block = self.fetchDRMdata()
  124.             else:
  125.                 raise MobiError('Unable to set metadata on DRM file without EXTH header')
  126.         self.have_exth
  127.         self.original_exth_records = { }
  128.         self.fetchEXTHFields()
  129.  
  130.     
  131.     def fetchDRMdata(self):
  132.         drm_offset = int(unpack('>I', self.record0[168:172])[0])
  133.         self.drm_key_count = int(unpack('>I', self.record0[172:176])[0])
  134.         drm_keys = ''
  135.         for x in range(self.drm_key_count):
  136.             base_addr = drm_offset + x * self.DRM_KEY_SIZE
  137.             drm_keys += self.record0[base_addr:base_addr + self.DRM_KEY_SIZE]
  138.         
  139.         return drm_keys
  140.  
  141.     
  142.     def fetchEXTHFields(self):
  143.         stream = self.stream
  144.         record0 = self.record0
  145.         exth_off = unpack('>I', record0[20:24])[0] + 16 + record0.start
  146.         (image_base,) = unpack('>I', record0[108:112])
  147.         exth = self.exth = StreamSlicer(stream, exth_off, record0.stop)
  148.         (nitems,) = unpack('>I', exth[8:12])
  149.         pos = 12
  150.         for i in xrange(nitems):
  151.             (id, size) = unpack('>II', exth[pos:pos + 8])
  152.             content = exth[pos + 8:pos + size]
  153.             pos += size
  154.             self.original_exth_records[id] = content
  155.             if id == 106:
  156.                 self.timestamp = content
  157.                 continue
  158.             if id == 201:
  159.                 (rindex,) = (self.cover_rindex,) = unpack('>i', content)
  160.                 if rindex > 0:
  161.                     self.cover_record = self.record(rindex + image_base)
  162.                 
  163.             rindex > 0
  164.             if id == 202:
  165.                 (rindex,) = (self.thumbnail_rindex,) = unpack('>i', content)
  166.                 if rindex > 0:
  167.                     self.thumbnail_record = self.record(rindex + image_base)
  168.                 
  169.             rindex > 0
  170.         
  171.  
  172.     
  173.     def patch(self, off, new_record0):
  174.         record_sizes = [
  175.             len(new_record0)]
  176.         for i in range(1, self.nrecs - 1):
  177.             record_sizes.append(self.pdbrecords[i + 1][0] - self.pdbrecords[i][0])
  178.         
  179.         record_sizes.append(self.data.stop - self.pdbrecords[self.nrecs - 1][0])
  180.         updated_pdbrecords = [
  181.             self.pdbrecords[0][0]]
  182.         record0_offset = self.pdbrecords[0][0]
  183.         updated_offset = record0_offset + len(new_record0)
  184.         for i in range(1, self.nrecs - 1):
  185.             updated_pdbrecords.append(updated_offset)
  186.             updated_offset += record_sizes[i]
  187.         
  188.         updated_pdbrecords.append(updated_offset)
  189.         data_blocks = [
  190.             new_record0]
  191.         for i in range(1, self.nrecs):
  192.             data_blocks.append(self.data[self.pdbrecords[i][0]:self.pdbrecords[i][0] + record_sizes[i]])
  193.         
  194.         self.record0.update(data_blocks)
  195.         self.update_pdbrecords(updated_pdbrecords)
  196.         if updated_pdbrecords[-1] + record_sizes[-1] < self.data.stop:
  197.             self.data.truncate(updated_pdbrecords[-1] + record_sizes[-1])
  198.         else:
  199.             self.data.stop = updated_pdbrecords[-1] + record_sizes[-1]
  200.  
  201.     
  202.     def patchSection(self, section, new):
  203.         off = self.pdbrecords[section][0]
  204.         self.patch(off, new)
  205.  
  206.     
  207.     def create_exth(self, new_title = None, exth = None):
  208.         if isinstance(new_title, unicode):
  209.             new_title = new_title.encode(self.codec, 'replace')
  210.         
  211.         (title_offset,) = unpack('>L', self.record0[84:88])
  212.         (title_length,) = unpack('>L', self.record0[88:92])
  213.         (title_in_file,) = unpack('%ds' % title_length, self.record0[title_offset:title_offset + title_length])
  214.         (mobi_header_length,) = unpack('>L', self.record0[20:24])
  215.         if mobi_header_length == 228:
  216.             self.record0[23] = '\xe8'
  217.             self.record0[244:248] = pack('>L', 0xFFFFFFFFL)
  218.             mobi_header_length = 232
  219.         
  220.         self.record0[128:132] = pack('>L', self.flags | 64)
  221.         if not exth:
  222.             pad = '\x00\x00\x00\x00'
  223.             exth = [
  224.                 'EXTH',
  225.                 pack('>II', 12, 0),
  226.                 pad]
  227.             exth = ''.join(exth)
  228.         
  229.         if self.encryption_type != 0:
  230.             self.record0[168:172] = pack('>L', 16 + mobi_header_length + len(exth))
  231.             self.record0[176:180] = pack('>L', len(self.drm_block))
  232.             self.record0[84:88] = pack('>L', 16 + mobi_header_length + len(exth) + len(self.drm_block))
  233.         else:
  234.             self.record0[84:88] = pack('>L', 16 + mobi_header_length + len(exth))
  235.         if new_title:
  236.             self.record0[88:92] = pack('>L', len(new_title))
  237.         
  238.         new_record0 = StringIO()
  239.         new_record0.write(self.record0[:16 + mobi_header_length])
  240.         new_record0.write(exth)
  241.         if self.encryption_type != 0:
  242.             new_record0.write(self.drm_block)
  243.         
  244.         None(new_record0.write if new_title else title_in_file)
  245.         trail = len(new_record0.getvalue()) % 4
  246.         pad = '\x00' * (4 - trail)
  247.         new_record0.write(pad)
  248.         self.patchSection(0, new_record0.getvalue())
  249.         self.record0 = self.record(0)
  250.  
  251.     
  252.     def hexdump(self, src, length = 16):
  253.         FILTER = []([ '.' for x in range(256) ])
  254.         N = 0
  255.         result = ''
  256.         for x in s:
  257.             hexa = _[2](_[2]['%02X' % ord(x)])
  258.             s = s.translate(FILTER)
  259.             result += '%04X   %-*s   %s\n' % (N, length * 3, hexa, s)
  260.             N += length
  261.             []
  262.         print result
  263.  
  264.     
  265.     def get_pdbrecords(self):
  266.         pdbrecords = []
  267.         for i in xrange(self.nrecs):
  268.             (offset, a1, a2, a3, a4) = unpack('>LBBBB', self.data[78 + i * 8:78 + i * 8 + 8])
  269.             flags = a1
  270.             val = a2 << 16 | a3 << 8 | a4
  271.             pdbrecords.append([
  272.                 offset,
  273.                 flags,
  274.                 val])
  275.         
  276.         return pdbrecords
  277.  
  278.     
  279.     def update_pdbrecords(self, updated_pdbrecords):
  280.         for i, pdbrecord in enumerate(updated_pdbrecords):
  281.             self.data[78 + i * 8:78 + i * 8 + 4] = pack('>L', pdbrecord)
  282.         
  283.         self.pdbrecords = self.get_pdbrecords()
  284.  
  285.     
  286.     def dump_pdbrecords(self):
  287.         print 'MetadataUpdater.dump_pdbrecords()'
  288.         print '%10s %10s %10s' % ('offset', 'flags', 'val')
  289.         for i in xrange(len(self.pdbrecords)):
  290.             pdbrecord = self.pdbrecords[i]
  291.             print '%10X %10X %10X' % (pdbrecord[0], pdbrecord[1], pdbrecord[2])
  292.         
  293.  
  294.     
  295.     def record(self, n):
  296.         if n >= self.nrecs:
  297.             raise ValueError('non-existent record %r' % n)
  298.         n >= self.nrecs
  299.         offoff = 78 + 8 * n
  300.         (start,) = unpack('>I', self.data[offoff + 0:offoff + 4])
  301.         stop = None
  302.         if n < self.nrecs - 1:
  303.             (stop,) = unpack('>I', self.data[offoff + 8:offoff + 12])
  304.         
  305.         return StreamSlicer(self.stream, start, stop)
  306.  
  307.     
  308.     def update(self, mi):
  309.         
  310.         def update_exth_record(rec):
  311.             recs.append(rec)
  312.             if rec[0] in self.original_exth_records:
  313.                 self.original_exth_records.pop(rec[0])
  314.             
  315.  
  316.         if self.type != 'BOOKMOBI':
  317.             raise MobiError("Setting metadata only supported for MOBI files of type 'BOOK'.\n\tThis is a '%s' file of type '%s'" % (self.type[0:4], self.type[4:8]))
  318.         self.type != 'BOOKMOBI'
  319.         recs = []
  320.         
  321.         try:
  322.             load_defaults = load_defaults
  323.             import calibre.ebooks.conversion.config
  324.             prefs = load_defaults('mobi_output')
  325.             pas = prefs.get('prefer_author_sort', False)
  326.             kindle_pdoc = prefs.get('personal_doc', None)
  327.         except:
  328.             (None, None)
  329.             pas = False
  330.             kindle_pdoc = None
  331.  
  332.         if mi.author_sort and pas:
  333.             authors = mi.author_sort
  334.             update_exth_record((100, authors.encode(self.codec, 'replace')))
  335.         elif mi.authors:
  336.             authors = ';'.join(mi.authors)
  337.             update_exth_record((100, authors.encode(self.codec, 'replace')))
  338.         
  339.         if mi.publisher:
  340.             update_exth_record((101, mi.publisher.encode(self.codec, 'replace')))
  341.         
  342.         if mi.comments:
  343.             a_offset = mi.comments.find('<div class="user_annotations">')
  344.             ad_offset = mi.comments.find('<hr class="annotations_divider" />')
  345.             if a_offset >= 0:
  346.                 mi.comments = mi.comments[:a_offset]
  347.             
  348.             if ad_offset >= 0:
  349.                 mi.comments = mi.comments[:ad_offset]
  350.             
  351.             update_exth_record((103, mi.comments.encode(self.codec, 'replace')))
  352.         
  353.         if mi.isbn:
  354.             update_exth_record((104, mi.isbn.encode(self.codec, 'replace')))
  355.         
  356.         if mi.tags:
  357.             subjects = '; '.join(mi.tags)
  358.             update_exth_record((105, subjects.encode(self.codec, 'replace')))
  359.             if kindle_pdoc and kindle_pdoc in mi.tags:
  360.                 update_exth_record((501, str('PDOC')))
  361.             
  362.         
  363.         if mi.pubdate:
  364.             update_exth_record((106, str(mi.pubdate).encode(self.codec, 'replace')))
  365.         elif mi.timestamp:
  366.             update_exth_record((106, str(mi.timestamp).encode(self.codec, 'replace')))
  367.         elif self.timestamp:
  368.             update_exth_record((106, self.timestamp))
  369.         else:
  370.             update_exth_record((106, nowf().isoformat().encode(self.codec, 'replace')))
  371.         if self.cover_record is not None:
  372.             update_exth_record((201, pack('>I', self.cover_rindex)))
  373.             update_exth_record((203, pack('>I', 0)))
  374.         
  375.         if self.thumbnail_record is not None:
  376.             update_exth_record((202, pack('>I', self.thumbnail_rindex)))
  377.         
  378.         if 503 in self.original_exth_records:
  379.             update_exth_record((503, mi.title.encode(self.codec, 'replace')))
  380.         
  381.         for id in sorted(self.original_exth_records):
  382.             recs.append((id, self.original_exth_records[id]))
  383.         
  384.         recs = sorted(recs, key = (lambda x: (x[0], x[0])))
  385.         exth = StringIO()
  386.         for code, data in recs:
  387.             exth.write(pack('>II', code, len(data) + 8))
  388.             exth.write(data)
  389.         
  390.         exth = exth.getvalue()
  391.         trail = len(exth) % 4
  392.         pad = '\x00' * (4 - trail)
  393.         exth = [
  394.             'EXTH',
  395.             pack('>II', len(exth) + 12, len(recs)),
  396.             exth,
  397.             pad]
  398.         exth = ''.join(exth)
  399.         if getattr(self, 'exth', None) is None:
  400.             raise MobiError('No existing EXTH record. Cannot update metadata.')
  401.         getattr(self, 'exth', None) is None
  402.         self.record0[92:96] = iana2mobi(mi.language)
  403.         self.create_exth(exth = exth, new_title = mi.title)
  404.         self.fetchEXTHFields()
  405.         if mi.cover_data[1] or mi.cover:
  406.             
  407.             try:
  408.                 data = None if mi.cover_data[1] else open(mi.cover, 'rb').read()
  409.             except:
  410.                 pass
  411.  
  412.             if self.cover_record is not None:
  413.                 size = len(self.cover_record)
  414.                 cover = rescale_image(data, size)
  415.                 cover += '\x00' * (size - len(cover))
  416.                 self.cover_record[:] = cover
  417.             
  418.             if self.thumbnail_record is not None:
  419.                 size = len(self.thumbnail_record)
  420.                 thumbnail = rescale_image(data, size, dimen = MAX_THUMB_DIMEN)
  421.                 thumbnail += '\x00' * (size - len(thumbnail))
  422.                 self.thumbnail_record[:] = thumbnail
  423.             
  424.         
  425.  
  426.  
  427.  
  428. def set_metadata(stream, mi):
  429.     mu = MetadataUpdater(stream)
  430.     mu.update(mi)
  431.  
  432.