Obok - Kobo DRM removal

#!/usr/bin/python
##########################################################
#                    KOBO DRM CRACK BY                   #
#                      PHYSISTICATED                     #
#                      8====D~~~~ O:                     #
##########################################################
# This app was made for Python 2.7 on Windows 32-bit
#
# This app needs pycrypto - get from here:
# http://www.voidspace.org.uk/python/modules.shtml
#
# Usage: obok.py
# Choose the book you want to decrypt
#
# Shouts to my krew - you know who you are - and one in
# particular who gave me a lot of help with this - thank
# you so much!
#
# Kopimi /K\
# Keep sharing, keep copying, but remember that nothing is
# for free - make sure you compensate your favorite
# authors - and cut out the middle man whenever possible
# ;) ;) ;)
#
# DRM AUTOPSY
# The Kobo DRM was incredibly easy to crack, but it took
# me months to get around to making this. Here's the
# basics of how it works:
# 1: Get MAC address of first NIC in ipconfig (sometimes
# stored in registry as pwsdid)
# 2: Get user ID (stored in tons of places, this gets it
# from HKEY_CURRENT_USER\Software\Kobo\Kobo Desktop
# Edition\Browser\cookies)
# 3: Concatenate and SHA256, take the second half - this
# is your master key
# 4: Open %LOCALAPPDATA%\Kobo Desktop Editions\Kobo.sqlite
# and dump content_keys
# 5: Unbase64 the keys, then decode these with the master
# key - these are your page keys
# 6: Unzip EPUB of your choice, decrypt each page with its
# page key, then zip back up again
#
# WHY USE THIS WHEN INEPT WORKS FINE? (adobe DRM stripper)
# Inept works very well, but authors on Kobo can choose
# what DRM they want to use - and some have chosen not to
# let people download them with Adobe Digital Editions -
# they would rather lock you into a single platform.
#
# With Obok, you can sync Kobo Desktop, decrypt all your
# ebooks, and then use them on whatever device you want
# - you bought them, you own them, you can do what you
# like with them.
#
# Obok is Kobo backwards, but it is also means "next to"
# in Polish.
# When you buy a real book, it is right next to you. You
# can read it at home, at work, on a train, you can lend
# it to a friend, you can scribble on it, and add your own
# explanations/translations.
#
# Obok gives you this power over your ebooks - no longer
# are you restricted to one device. This allows you to
# embed foreign fonts into your books, as older Kobo's
# can't display them properly. You can read your books
# on your phones, in different PC readers, and different
# ereader devices. You can share them with your friends
# too, if you like - you can do that with a real book
# after all.
#

import os
import _winreg
import re
import string
import hashlib
import sqlite3
import base64
import binascii
import zipfile
from Crypto.Cipher import AES
from uuid import getnode as get_mac

def SHA256(raw):
    return hashlib.sha256(raw).hexdigest()

def RemoveAESPadding(contents):
    lastchar = binascii.b2a_hex(contents[-1:])
    len = int(lastchar, 16)
    padding = len
    if(len == 1):
        return contents[:-1]
    if(len < 16):
        for i in range(len):
            testchar = binascii.b2a_hex(contents[-len:-(len-1)])
            if(testchar != lastchar):
                padding = 0
    if(padding > 0):
        contents = contents[:-padding]
    return contents

def GetVolumeKeys(c):
    volumekeys = {}
    for row in c.execute("SELECT * from content_keys"):
        if(row[0] not in volumekeys):
            volumekeys[row[0]] = {}
        volumekeys[row[0]][row[1]] = {}
        volumekeys[row[0]][row[1]]["encryptedkey"] = base64.b64decode(row[2])
        volumekeys[row[0]][row[1]]["decryptedkey"] = enc.decrypt(volumekeys[row[0]][row[1]]["encryptedkey"])
    # get book name
    for key in volumekeys.keys():
        volumekeys[key]["title"] = c.execute("SELECT Title from content where ContentID = '%s'" % (key)).fetchone()[0]
    return volumekeys


def GetMacString():
    # terrible way, but get_mac() only returns proper interfaces
    # this will return the first in ipconfig
    # Kobo appears to do the same
    ipconfig = os.popen('ipconfig /all').read()
    c = re.compile("Physical Address[^0-9A-F]+([0-9A-Fa-f-]+)", re.MULTILINE)
    m = c.search(ipconfig)
    mac = re.sub("-", ":", m.group(1))
    return mac

    #macnum = get_mac()
    #macarray = []
    #for i in range(6):
    #   macarray.append((hex((macnum >> (40-(i*8)) & 0xFF))[2:-1]).zfill(2))
    #return string.upper(string.join(macarray, ":"))

def ByteArrayToString(bytearray):
    return re.match("@ByteArray\\((.+)\\)", bytearray).group(1)

def GetUserHexKey():
    regkey_browser = _winreg.OpenKey(_winreg.HKEY_CURRENT_USER, "Software\\Kobo\\Kobo Desktop Edition\\Browser")
    cookies = _winreg.QueryValueEx(regkey_browser, "cookies")
    bytearrays = cookies[0]
    # find wsuid and pwsdid
    # pwsdid isn't always stored, but GetMacString() will get it for us
    wsuid = ""
    pwsdid = GetMacString()
    for bytearray in bytearrays:
        cookie = ByteArrayToString(bytearray)
        wsuidcheck = re.match("^wsuid=([0-9a-f-]+)", cookie)
        if(wsuidcheck):
            wsuid = wsuidcheck.group(1)
    if(wsuid == ""):
        print "wsuid key not found :/"
        exit()
    preuserkey = string.join((pwsdid, wsuid), "")
    userkey = SHA256(preuserkey)
    return userkey[32:]

# get key
userkeyhex = GetUserHexKey()
# load into AES
userkey = binascii.a2b_hex(userkeyhex)
enc = AES.new(userkey, AES.MODE_ECB)
# get dirs
kobodir = string.join((os.environ['LOCALAPPDATA'], "Kobo\\Kobo Desktop Edition"), "\\")
sqlitefile = string.join((kobodir, "Kobo.sqlite"), "\\")
bookdir = string.join((kobodir, "kepub"), "\\")

# open sqlite
conn = sqlite3.connect(sqlitefile)
c = conn.cursor()
# get volume keys
volumekeys = GetVolumeKeys(c)

# choose a volumeID

volumeid = ""
print "Choose a book to decrypt:"
i = 1
for key in volumekeys.keys():
    print "%d: %s" % (i, volumekeys[key]["title"])
    i+=1

num = input("...")

i = 1
for key in volumekeys.keys():
    if(i==num):
        volumeid = key
    i+=1


if(volumeid == ""):
    exit()

zippath = string.join((bookdir, volumeid), "\\")

z = zipfile.ZipFile(zippath, "r")
# make filename out of a-zA-Z0-9 from title
outname = "%s.epub" % (re.sub("[^a-zA-Z0-9 ]", "", volumekeys[volumeid]["title"]))
zout = zipfile.ZipFile(outname, "w", zipfile.ZIP_DEFLATED)
for filename in z.namelist():
    #print filename
    # read in and decrypt
    if(filename in volumekeys[volumeid]):
        # do decrypted version
        pagekey = volumekeys[volumeid][filename]["decryptedkey"]
        penc = AES.new(pagekey, AES.MODE_ECB)
        contents = RemoveAESPadding(penc.decrypt(z.read(filename)))
        # need to fix padding
        zout.writestr(filename, contents)
    else:
        zout.writestr(filename, z.read(filename))

print "Book saved as %s\\%s" % (os.getcwd(), outname)