#!/usr/bin/python
#
# Changelog
#  1.00 - Initial version, with code from various other scripts.
#  1.01 - Moved authorship announcement to usage section.
# Written in 2011 by Paul Durrant
#
#  1.02 - Added recognition of Apple Fairplay encryption.
# Modified in 2015 by Alasdair Allan
#
# Released with unlicense. See http://unlicense.org/
#
#############################################################################
#
# This is free and unencumbered software released into the public domain.
# 
# Anyone is free to copy, modify, publish, use, compile, sell, or
# distribute this software, either in source code form or as a compiled
# binary, for any purpose, commercial or non-commercial, and by any
# means.
# 
# In jurisdictions that recognize copyright laws, the author or authors
# of this software dedicate any and all copyright interest in the
# software to the public domain. We make this dedication for the benefit
# of the public at large and to the detriment of our heirs and
# successors. We intend this dedication to be an overt act of
# relinquishment in perpetuity of all present and future rights to this
# software under copyright law.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
#############################################################################

from __future__ import with_statement

__version__ = '1.02'

import sys, struct, os
import zlib
import zipfile

_FILENAME_LEN_OFFSET = 26
_EXTRA_LEN_OFFSET = 28
_FILENAME_OFFSET = 30
_MAX_SIZE = 64 * 1024


def uncompress(cmpdata):
	dc = zlib.decompressobj(-15)
	data = ''
	while len(cmpdata) > 0:
		if len(cmpdata) > _MAX_SIZE :
			newdata = cmpdata[0:_MAX_SIZE]
			cmpdata = cmpdata[_MAX_SIZE:]
		else:
			newdata = cmpdata
			cmpdata = ''
		newdata = dc.decompress(newdata)
		unprocessed = dc.unconsumed_tail
		if len(unprocessed) == 0:
			newdata += dc.flush()
		data += newdata
		cmpdata += unprocessed
		unprocessed = ''
	return data

def getfiledata(file, zi):
	# get file name length and exta data length to find start of file data
	local_header_offset = zi.header_offset

	file.seek(local_header_offset + _FILENAME_LEN_OFFSET)
	leninfo = file.read(2)
	local_name_length, = struct.unpack('<H', leninfo)

	file.seek(local_header_offset + _EXTRA_LEN_OFFSET)
	exinfo = file.read(2)
	extra_field_length, = struct.unpack('<H', exinfo)

	file.seek(local_header_offset + _FILENAME_OFFSET + local_name_length + extra_field_length)
	data = None

	# if not compressed we are good to go
	if zi.compress_type == zipfile.ZIP_STORED:
		data = file.read(zi.file_size)

	# if compressed we must decompress it using zlib
	if zi.compress_type == zipfile.ZIP_DEFLATED:
		cmpdata = file.read(zi.compress_size)
		data = uncompress(cmpdata)

	return data

def main():
	if len(sys.argv) < 2:
		print "drmcheck v%s    Written 2011 Paul Durrant" % __version__
		print "Released with unlicense. See http://unlicense.org/\n"
		print "\nDescription:\n  Determines whether given ebook has DRM."
		print "Usage:\n  %s <infile>" % sys.argv[0]
		return 1
	infile = sys.argv[1]
	kind = "Unknown"
	compression = ""
	encryption = ""
	infileobject = file(infile,'rb')
	bookdata = infileobject.read()
	# Check for Mobipocket/Kindle
	if bookdata[60:60+8] == 'BOOKMOBI':
		kind = "Mobipocket"
		offset, = struct.unpack('>L',bookdata[78:78+4])
		compressionid, = struct.unpack('>H',bookdata[offset:offset+2])
		if compressionid==1:
			compression = "uncompressed"
		elif compressionid==2:
			compression = "PalmDOC compression"
		elif compressionid==17480:
			compression = "HUFF/CDIC compression"
		else:
			compression = "unknown compression type %d" % compresionid
		encryptionid, = struct.unpack('>H',bookdata[offset+12:offset+12+2])
		if encryptionid==0:
			encryption = "unencrypted"
		elif encryptionid==1:
			encryption = "encrypted with old encrytion method"
		elif encryptionid==2:
			encryption = "encrypted with current encrytion method"
		else:
			encryption = "encrypted with unknown encryption method %d" % encryptionid
	if bookdata[60:60+8] == 'TEXtREAd':
		kind = "PalmDoc or early Mobipocket"
		offset, = struct.unpack('>L',bookdata[78:78+4])
		compressionid, = struct.unpack('>H',bookdata[offset:offset+2])
		if compressionid==1:
			compression = "uncompressed"
		elif compressionid==2:
			compression = "PalmDOC compression"
		else:
			compression = "unknown compression type %d" % compresionid
		encryptionid, = struct.unpack('>H',bookdata[offset+12:offset+12+2])
		if encryptionid==1:
			encryption = "encrypted with old encrytion method"
		elif encryptionid==2:
			encryption = "encrypted with current encrytion method"
		else:
			encryption = "unencrypted"
	if bookdata[60:60+8] == 'PNRdPPrs':
		kind = "eReader"
		offset, = struct.unpack('>L',bookdata[78:78+4])
		compressionid, = struct.unpack('>H',bookdata[offset:offset+2])
		if compressionid==2:
			compression = "PalmDOC compression"
			encryption = "unencrypted"
		elif compressionid==10:
			compression = "zlib compression"
			encryption = "unencrypted"
		elif compressionid==260 or compressionid==259 or compressionid==272:
			compression = "zlib compression"
			encryption = "encrypted"
		else:
			compression = "unknown compression/encryption type %d" % compresionid
			encryption = ""
	if bookdata[0:0+4] == 'TPZ0':
		kind = "Amazon Topaz"
		compression = "compressed"
		encryption = "encrypted"
	if bookdata[0:0+2] == "PK":
		if bookdata[30:30+28] == 'mimetypeapplication/epub+zip':
			kind = "ePub"
		else:
			kind = "Probably a ZIP file, possibly an ePub"
		compression = ""
		encryption = "unencrypted"
		foundrights = 0
		foundencryption = 0
		inzip = zipfile.ZipFile(infile,'r')
		for zinfo in inzip.infolist():
			if zinfo.filename.find("encryption.xml") != -1:
				foundencryption = 1
				if foundrights == 0:
					encryption = "Unknown encryption type"
			if zinfo.filename.find("sinf.xml") != -1:
				foundencryption = 1
				fairdata = getfiledata(infileobject, zinfo)
				offset = fairdata.find("<fairplay:sID>")
				if fairdata[offset+14] == "1":
					encryption = "Apple Fairplay"
			if zinfo.filename.find("rights.xml") != -1:
				foundrights = 1
				encryption = "encrypted"
				rightsdata = getfiledata(infileobject, zinfo)
				firstoffset = rightsdata.find("<encryptedKey>")
				secondoffset = rightsdata.find("</encryptedKey>")
				if secondoffset-firstoffset == 78:
					encryption = "Barnes & Noble encryption"
				elif secondoffset-firstoffset == 186:
					encryption = "Adobe Digital Editions encryption"
				else:
					encryption = "Unknown encryption type"
		if foundencryption == 0:
			encryption = "unencrypted"
	infileobject.close()
		
	print "%s - %s ebook, %s %s.\n" % (infile, kind, compression, encryption)
	return 0

if __name__ == "__main__":
	sys.exit(main())