Preliminary implementation of DualMetaFix

1 files changed, 183 insertions, 0 deletions

diff --git a/kcc/dualmetafix.py b/kcc/dualmetafix.py
new file mode 100644
index 0000000..db07933
--- /dev/null
+++ b/kcc/dualmetafix.py
@@ -0,0 +1,183 @@
+# -*- coding: utf-8 -*-
+#
+# Based on initial version of DualMetaFix. Copyright (C) 2013 Kevin Hendricks
+# Changes for KCC Copyright (C) 2014 Pawel Jastrzebski <pawelj@vulturis.eu>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import struct
+
+
+class DualMetaFixException(Exception):
+    pass
+
+# palm database offset constants
+number_of_pdb_records = 76
+first_pdb_record = 78
+
+# important rec0 offsets
+mobi_header_base = 16
+mobi_header_length = 20
+mobi_version = 36
+title_offset = 84
+
+
+def getint(data, ofs, sz='L'):
+    i, = struct.unpack_from('>'+sz, data, ofs)
+    return i
+
+
+def writeint(data, ofs, n, slen='L'):
+    if slen == 'L':
+        return data[:ofs]+struct.pack('>L', n)+data[ofs+4:]
+    else:
+        return data[:ofs]+struct.pack('>H', n)+data[ofs+2:]
+
+
+def getsecaddr(datain, secno):
+    nsec = getint(datain, number_of_pdb_records, 'H')
+    if (secno < 0) | (secno >= nsec):
+        emsg = 'requested section number %d out of range (nsec=%d)' % (secno, nsec)
+        raise DualMetaFixException(emsg)
+    secstart = getint(datain, first_pdb_record+secno*8)
+    if secno == nsec-1:
+        secend = len(datain)
+    else:
+        secend = getint(datain, first_pdb_record+(secno+1)*8)
+    return secstart, secend
+
+
+def readsection(datain, secno):
+    secstart, secend = getsecaddr(datain, secno)
+    return datain[secstart:secend]
+
+
+# overwrite section - must be exact same length as original
+def replacesection(datain, secno, secdata):
+    secstart, secend = getsecaddr(datain, secno)
+    seclen = secend - secstart
+    if len(secdata) != seclen:
+        raise DualMetaFixException('section length change in replacesection')
+    datalst = [datain[0:secstart], secdata, datain[secend:]]
+    dataout = b"".join(datalst)
+    return dataout
+
+
+def get_exth_params(rec0):
+    ebase = mobi_header_base + getint(rec0, mobi_header_length)
+    if rec0[ebase:ebase+4] != b'EXTH':
+        raise DualMetaFixException('EXTH tag not found where expected')
+    elen = getint(rec0, ebase+4)
+    enum = getint(rec0, ebase+8)
+    rlen = len(rec0)
+    return ebase, elen, enum, rlen
+
+
+def add_exth(rec0, exth_num, exth_bytes):
+    ebase, elen, enum, rlen = get_exth_params(rec0)
+    newrecsize = 8+len(exth_bytes)
+    newrec0 = rec0[0:ebase+4]+struct.pack('>L', elen+newrecsize)+struct.pack('>L', enum+1)+struct.pack('>L', exth_num)\
+        + struct.pack('>L', newrecsize)+exth_bytes+rec0[ebase+12:]
+    newrec0 = writeint(newrec0, title_offset, getint(newrec0, title_offset)+newrecsize)
+    # keep constant record length by removing newrecsize null bytes from end
+    sectail = newrec0[-newrecsize:]
+    if sectail != b'\0'*newrecsize:
+        raise DualMetaFixException('add_exth: trimmed non-null bytes at end of section')
+    newrec0 = newrec0[0:rlen]
+    return newrec0
+
+
+def read_exth(rec0, exth_num):
+    exth_values = []
+    ebase, elen, enum, rlen = get_exth_params(rec0)
+    ebase += 12
+    while enum > 0:
+        exth_id = getint(rec0, ebase)
+        if exth_id == exth_num:
+            # We might have multiple exths, so build a list.
+            exth_values.append(rec0[ebase+8:ebase+getint(rec0, ebase+4)])
+        enum -= 1
+        ebase = ebase+getint(rec0, ebase+4)
+    return exth_values
+
+
+def del_exth(rec0, exth_num):
+    ebase, elen, enum, rlen = get_exth_params(rec0)
+    ebase_idx = ebase+12
+    enum_idx = 0
+    while enum_idx < enum:
+        exth_id = getint(rec0, ebase_idx)
+        exth_size = getint(rec0, ebase_idx+4)
+        if exth_id == exth_num:
+            newrec0 = rec0
+            newrec0 = writeint(newrec0, title_offset, getint(newrec0, title_offset)-exth_size)
+            newrec0 = newrec0[:ebase_idx]+newrec0[ebase_idx+exth_size:]
+            newrec0 = newrec0[0:ebase+4]+struct.pack('>L', elen-exth_size)+struct.pack('>L', enum-1)+newrec0[ebase+12:]
+            newrec0 += b'\0'*exth_size
+            if rlen != len(newrec0):
+                raise DualMetaFixException('del_exth: incorrect section size change')
+            return newrec0
+        enum_idx += 1
+        ebase_idx = ebase_idx+exth_size
+    return rec0
+
+
+class DualMobiMetaFix:
+
+    def __init__(self, infile, asin):
+        self.datain = open(infile, 'rb').read()
+        self.datain_rec0 = readsection(self.datain, 0)
+
+        # in the first mobi header
+        # add 501 to "EBOK", add 113 as asin, add 504 as asin
+        rec0 = self.datain_rec0
+        rec0 = del_exth(rec0, 501)
+        rec0 = del_exth(rec0, 113)
+        rec0 = del_exth(rec0, 504)
+        rec0 = add_exth(rec0, 501, b"EBOK")
+        rec0 = add_exth(rec0, 113, asin)
+        rec0 = add_exth(rec0, 504, asin)
+        self.datain = replacesection(self.datain, 0, rec0)
+
+        ver = getint(self.datain_rec0, mobi_version)
+        self.combo = (ver != 8)
+        if not self.combo:
+            return
+
+        exth121 = read_exth(self.datain_rec0, 121)
+        if len(exth121) == 0:
+            self.combo = False
+            return
+        else:
+            # only pay attention to first exth121
+            # (there should only be one)
+            datain_kf8, = struct.unpack_from('>L', exth121[0], 0)
+            if datain_kf8 == 0xffffffff:
+                self.combo = False
+                return
+        self.datain_kfrec0 = readsection(self.datain, datain_kf8)
+
+        # in the second header
+        # add 501 to "EBOK", add 113 as asin, add 504 as asin
+        rec0 = self.datain_kfrec0
+        rec0 = del_exth(rec0, 501)
+        rec0 = del_exth(rec0, 113)
+        rec0 = del_exth(rec0, 504)
+        rec0 = add_exth(rec0, 501, b"EBOK")
+        rec0 = add_exth(rec0, 113, asin)
+        rec0 = add_exth(rec0, 504, asin)
+        self.datain = replacesection(self.datain, datain_kf8, rec0)
+
+    def getresult(self):
+        return self.datain
\ No newline at end of file