from pwn import *

class PaddingOracle():
	def __init__(self, key_size=16):
		if key_size%8!=0:
			raise ValueError("Incorrect key length: %d",key_size)
		self.key_size=key_size

	def oracle(self, payload, iv, previous_resp, **kwargs):
		"""
			payload - hex encoded ciphertext to check
			iv - hex encoded inicialization vector (most often: append it at beggining of payload)
			previous_resp - value returned by previous call to oracle
		Return tuple: (True/False, response)
		True if padding is correct, Flase otherwise; Be carefull while sending original ciphertext
		response will be send as previous_response to next call
		"""
		raise NotImplementedError

	def decrypt(self, ciphertext, iv=None, amount=0, is_correct=False, known_plaintext=None, async=0, **kwargs):
		"""
			ciphertext - hex encoded
			iv - usually it's first block of ciphertext, but may be specified separate
			amount - how much blocks decrypt (counting from last), default is all
			is_correct - if True, then first decrypted byte will be treated as correct padding (speed up decryption)
			known_plaintext - hex encoded, with padding, from end
			async - make asynchronous calls to oracle (not implemented yet)
			kwargs - push forward to oracle function
		Return hex encoded plaintext
		"""

		log.info("Start decrypt")
		log.debug(self.chunks(ciphertext, size=self.key_size*2))
		ciphertext=ciphertext.decode('hex')
		if len(ciphertext)%self.key_size!=0:
			raise ValueError("Incorrect ciphertext length: %d"%(len(ciphertext)))

		#prepare blocks
		blocks=re.findall('.{%d}'%(self.key_size), ciphertext, re.DOTALL)
		resp=None
		if iv!=None:
			log.info("Set iv")
			iv=iv.decode('hex')
			blocks.insert(0,iv)
		if amount!=0:
			amount=len(blocks)-amount-1
		if amount<0 or amount>=len(blocks):
			raise ValueError("Incorrect amount of blocks to decrypt: %d (have to be [0,%d)"%(amount,len(blocks)))
		log.info("Will decrypt %d block(s)"%(len(blocks)-1-amount))

		#add known plaintext
		plaintext=''
		position_known=0
		if known_plaintext!=None:
			kp=known_plaintext.decode('hex')
			plaintext=kp
			blocks_decoded=len(kp)/self.key_size
			chars_decoded=len(kp)%self.key_size
			if blocks_decoded>=len(blocks):
				raise RuntimeError("Too long known plaintext (%d blocks)"%blocks_decoded)
			blocks=blocks[:len(blocks)-blocks_decoded]
			blocks[-2]=blocks[-2][:-chars_decoded]+xor(kp[:chars_decoded],blocks[-2][-chars_decoded:],chars_decoded+1)
			position_known=chars_decoded
			log.info("Have known plaintext, skip %d block(s) and %d bytes"%(blocks_decoded, chars_decoded))

			
		for count_block in xrange(len(blocks)-1, amount, -1):
			""" Every block """
			log.info("Block no. %d"%count_block)

			payload_prefix=''.join(blocks[:count_block-1])
			payload_modify=blocks[count_block-1]
			payload_decrypt=blocks[count_block]

			position=self.key_size-1-position_known
			position_known=0
			while position>=0:
				""" Every position in block """
				log.debug("Position: %d"%position)

				for one in xrange(256):
					modified=payload_modify[:position]+chr(one)+payload_modify[position+1:]
					payload=''.join([payload_prefix, modified, payload_decrypt])

					iv=payload[:self.key_size].encode('hex')
					payload=payload[self.key_size:].encode('hex')

					is_ok=False
					log.debug(self.chunks( (iv+payload).encode('hex'), size=self.key_size*2))
					correct, resp=self.oracle(payload=payload, iv=iv, previous_resp=resp, **kwargs)
					if correct:
						""" oracle return True """ 
						padding=self.key_size-position
						decrypted_char=chr(ord(payload_modify[position])^one^padding)

						if is_correct==True and count_block==len(blocks)-1 and position==self.key_size-1:
							""" Have original padding value, can skip some bytes """
							dc=ord(decrypted_char)
							log.info("Found padding value: %d", dc)
							if dc==0 or dc>self.key_size:
								raise RuntimeError("Found bad padding value (given ciphertext may not be correct)")

							plaintext=decrypted_char*dc
							payload_modify=payload_modify[:-dc]+xor(payload_modify[-dc:],dc,dc+1)
							position=position-dc+1
						else:
							payload_modify=payload_modify[:position]+xor(chr(one)+payload_modify[position+1:], padding, padding+1)
							plaintext=decrypted_char+plaintext

						is_ok=True
						log.debug("Got correct char(\\x%x)"%one)
						log.info("Plaintext: %s"%plaintext.encode('hex'))
						log.debug("Plaintext: %s"%plaintext)
						break
				position-=1
				if is_ok==False:
					raise RuntimeError("Can't find corrent padding (oracle function return False 256 times)")
		log.success("Decrypted: %s"%plaintext.encode('hex'))
		return plaintext.encode('hex')



	def fake_ciphertext(self, new_plaintext, orginal_ciphertext, iv=None, orginal_plaintext=None, is_correct=False, **kwargs):
		"""
			new_plaintext - hex encoded, with padding, length==len(orginal_ciphertext+iv-key_size),You want returned ciphertext deciphered to this
			orginal_ciphertext - hex encoded, must be same length as new_plaintext
			iv - usually it's first block of ciphertext, but may be specified separate
			orginal_plaintext - hex encoded, can be only len(key_size) or None
			is_correct - if True, then first decrypted byte will be treated as correct padding (speed up decryption)
			kwargs - push forward to oracle function
		Return hex encoded fake ciphertext that will decrypt to new_plaintext
		"""

		log.info("Start fake ciphertext")
		orginal_ciphertext=orginal_ciphertext.decode('hex')
		new_plaintext=new_plaintext.decode('hex')
		if len(orginal_ciphertext)%self.key_size!=0:
			raise ValueError("Incorrect orginal ciphertext length: %d"%(len(ciphertext)))
		if len(new_plaintext)%self.key_size!=0:
			raise ValueError("Incorrect new plaintext length: %d"%(len(new_plaintext)))
		

		#prepare blocks
		blocks=re.findall('.{%d}'%(self.key_size), orginal_ciphertext, re.DOTALL)
		new_pl_blocks=re.findall('.{%d}'%(self.key_size), new_plaintext, re.DOTALL)
		new_pl_blocks.insert(0,'')
		if iv!=None:
			log.info("Set iv")
			iv=iv.decode('hex')
			blocks.insert(0,iv)
		if len(new_pl_blocks)!=len(blocks):
			raise RuntimeError("Wrong new plaintext length(%d), should be %d"%(len(new_plaintext),self.key_size*(len(blocks)-1)))
		new_ct_blocks=list(blocks)

		#add known plaintext
		if orginal_plaintext!=None:
			orginal_plaintext=orginal_plaintext.decode('hex')
			if orginal_plaintext>self.key_size:
				log.info("Cut original plaintext: from %d to last %d bytes"%(len(orginal_plaintext), self.key_size))
				orginal_plaintext=orginal_plaintext[-self.key_size:]
			orginal_plaintext=orginal_plaintext.encode('hex')

		for count_block in xrange(len(blocks)-1, 0, -1):
			""" Every block, modify block[count_block-1] to set block[count_block] """
			log.indented('')
			log.info("Block no. %d"%count_block)
			if orginal_plaintext==None:
				orginal_plaintext=self.decrypt(''.join(new_ct_blocks[:count_block+1]).encode('hex'), amount=1, is_correct=is_correct, **kwargs)
			elif len(orginal_plaintext)<self.key_size:
				orginal_plaintext=self.decrypt(''.join(new_ct_blocks[:count_block+1]).encode('hex'), amount=1, is_correct=is_correct, known_plaintext=orginal_plaintext, **kwargs)

			log.info("Set block no. %d"%count_block)
			orginal_plaintext=orginal_plaintext.decode('hex')
			new_ct_blocks[count_block-1]=xor(blocks[count_block-1], orginal_plaintext, new_pl_blocks[count_block])
			orginal_plaintext=None
			is_correct=False

		fake=''.join(new_ct_blocks).encode('hex')
		log.success("Fake ciphertext: %s"%fake)
		return fake

	def chunks(self, data, size, delim='|'):
		return delim.join([x for x in re.findall('.{%d}'%size, data, re.DOTALL)])

	def add_padding(self, data):
		size=self.key_size-len(data)%self.key_size
		return data+chr(size)*size