from Crypto.Cipher import AES
from Crypto.Util import Padding

from typing import Dict
import base64
import binascii
import hashlib
import zipfile

# Based on obok.py by Physisticated.
class KoboDrmRemover:
	def __init__( self, deviceId: str, userId: str ):
		self.DeviceIdUserIdKey = KoboDrmRemover.__MakeDeviceIdUserIdKey( deviceId, userId )

	@staticmethod
	def __MakeDeviceIdUserIdKey( deviceId: str, userId: str ) -> bytes:
		deviceIdUserId = ( deviceId + userId ).encode()
		key = hashlib.sha256( deviceIdUserId ).hexdigest()
		return binascii.a2b_hex( key[ 32: ] )

	def __DecryptContents( self, contents: bytes, contentKeyBase64: str ) -> bytes:
		contentKey = base64.b64decode( contentKeyBase64 )
		keyAes = AES.new( self.DeviceIdUserIdKey, AES.MODE_ECB )
		decryptedContentKey = keyAes.decrypt( contentKey )

		contentAes = AES.new( decryptedContentKey, AES.MODE_ECB )
		decryptedContents = contentAes.decrypt( contents )
		return Padding.unpad( decryptedContents, AES.block_size, "pkcs7" )

	def RemoveDrm( self, inputPath: str, outputPath: str, contentKeys: Dict[ str, str ] ) -> None:
		with zipfile.ZipFile( inputPath, "r" ) as inputZip:
			with zipfile.ZipFile( outputPath, "w", zipfile.ZIP_DEFLATED ) as outputZip:
				for filename in inputZip.namelist():
					contents = inputZip.read( filename )
					contentKeyBase64 = contentKeys.get( filename, None )
					if contentKeyBase64 is not None:
						contents = self.__DecryptContents( contents, contentKeyBase64 )
					outputZip.writestr( filename, contents )