
# Kudos:
# Parts of this code was inspired by the following project by @rubin_mor
# https://github.com/morRubin/AzureADJoinedMachinePTC
# 

# TODO: code currently supports RSA+DH+SHA1 , add support for other mechanisms


import os
import datetime
import secrets
import platform

from unicrypto import hashlib
from asn1crypto import cms
from asn1crypto import algos
from asn1crypto import core
from asn1crypto import x509
from asn1crypto import keys

from oscrypto.keys import parse_pkcs12
from oscrypto.asymmetric import rsa_pkcs1v15_sign, load_private_key

from minikerberos.protocol.constants import NAME_TYPE, MESSAGE_TYPE, PaDataType
from minikerberos.protocol.encryption import Enctype, _checksum_table, _enctype_table, Key
from minikerberos.protocol.structures import AuthenticatorChecksum
from minikerberos.protocol.asn1_structs import KDC_REQ_BODY, PrincipalName, HostAddress, \
	KDCOptions, EncASRepPart, AP_REQ, AuthorizationData, Checksum, krb5_pvno, Realm, \
	EncryptionKey, Authenticator, Ticket, APOptions, EncryptedData, AS_REQ, AP_REP
from minikerberos.protocol.rfc4556 import PKAuthenticator, AuthPack, Dunno1, Dunno2, MetaData, Info, CertIssuer, CertIssuers, PA_PK_AS_REP, KDCDHKeyInfo
from minikerberos.protocol.rfc_iakerb import KRB_FINISHED
from minikerberos.protocol.mskile import LSAP_TOKEN_INFO_INTEGRITY, KERB_AD_RESTRICTION_ENTRY, KERB_AD_RESTRICTION_ENTRYS
from minikerberos.gssapi.gssapi import GSSAPIFlags

def length_encode(x):
	if x <= 127:
		return x.to_bytes(1, 'big', signed = False)
	else:
		lb = x.to_bytes((x.bit_length() + 7) // 8, 'big')
		t = (0x80 | len(lb)).to_bytes(1, 'big', signed = False)
		return t+lb


class DirtyDH:
	def __init__(self):
		self.p = None
		self.g = None
		self.shared_key = None
		self.shared_key_int = None
		self.private_key = os.urandom(32)
		self.private_key_int = int(self.private_key.hex(), 16)
		self.dh_nonce = os.urandom(32)

	@staticmethod
	def from_params(p, g):
		dd = DirtyDH()
		dd.p = p
		dd.g = g
		return dd

	@staticmethod
	def from_dict(dhp):
		dd = DirtyDH()
		dd.p = dhp['p']
		dd.g = dhp['g']
		return dd

	@staticmethod
	def from_asn1(asn1_bytes):
		dhp = algos.DHParameters.load(asn1_bytes).native
		return DirtyDH.from_dict(dhp)
		
	
	def get_public_key(self):
		#y = g^x mod p
		return pow(self.g, self.private_key_int, self.p)
	
	def exchange(self, bob_int):
		self.shared_key_int = pow(bob_int, self.private_key_int, self.p)
		x = hex(self.shared_key_int)[2:]
		if len(x) % 2 != 0:
			x = '0' + x
		self.shared_key = bytes.fromhex(x)
		return self.shared_key


class PKINIT:
	def __init__(self):
		self.privkeyinfo = None
		self.certificate = None
		self.extra_certs = None
		self.user_sid = None
		self.user_name = None
		self.issuer = None
		self.cname = None
		self.diffie = None
		self.__hcert = None


	def init_windows_cert(self, username, certstore_name = 'MY', cert_serial = None):
		from minikerberos.common.windows.crypt32 import find_cert_by_cn
		self.certificate, self.__hcert = find_cert_by_cn(username, certstore_name)
		

	@staticmethod
	def from_windows_certstore(username, certstore_name = 'MY', cert_serial = None, dh_params = None):
		pkinit = PKINIT()
		pkinit.init_windows_cert(username, certstore_name = certstore_name, cert_serial = cert_serial)
		pkinit.setup(dh_params = dh_params)
		return pkinit


	@staticmethod
	def from_pfx(pfxfile, pfxpass, dh_params = None):
		pkinit = PKINIT()
		#print('Loading pfx12')		
		if isinstance(pfxpass, str):
			pfxpass = pfxpass.encode()
		with open(pfxfile, 'rb') as f:
			pkinit.privkeyinfo, pkinit.certificate, pkinit.extra_certs = parse_pkcs12(f.read(), password = pfxpass)
			pkinit.privkey = load_private_key(pkinit.privkeyinfo)
		#print('pfx12 loaded!')
		pkinit.setup(dh_params = dh_params)
		return pkinit
	
	def setup(self, dh_params = None):
		# parsing ceritficate to get basic info that will be needed to construct the asreq
		# this has two components
		for x in self.certificate.subject.native['common_name']:
			if x.startswith("S-1-12"):
				self.user_sid = x
			elif x.find('@') != -1:
				self.user_name = x

		self.issuer = self.certificate.issuer.native['common_name']
		self.cname = '\\'.join(['AzureAD', self.issuer, self.user_sid])

		#print('cert issuer: %s' % pkinit.issuer)
		#print('cert user_name: %s' % pkinit.user_name)
		#print('cert user_sid: %s' % pkinit.user_sid)
		#print('cert cname: %s' % pkinit.cname)

		if dh_params is None:
			print('Generating DH params...')
			self.diffie = DirtyDH.from_dict( generate_dh_parameters(1024).native)
			print('DH params generated.')
		else:
			#print('Loading default DH params...')
			if isinstance(dh_params, dict):
				self.diffie = DirtyDH.from_dict(dh_params)
			elif isinstance(dh_params, bytes):
				self.diffie = DirtyDH.from_asn1(dh_params)
			elif isinstance(dh_params, DirtyDH):
				self.diffie = dh_params
			else:
				raise Exception('DH params must be either a bytearray or a dict')


	def build_asreq(self, target = None, cname = None, kdcopts = ['forwardable','renewable','proxiable', 'canonicalize']):
		if isinstance(kdcopts, list):
			kdcopts = set(kdcopts)
		if cname is not None:
			if isinstance(cname, str):
				cname = [cname]
		else:
			cname = [self.cname]
		
		if target is not None:
			if isinstance(target, str):
				target = [target]
		else:
			target = ['127.0.0.1']

		now = datetime.datetime.now(datetime.timezone.utc)

		kdc_req_body_data = {}
		kdc_req_body_data['kdc-options'] = KDCOptions(kdcopts)
		kdc_req_body_data['cname'] = PrincipalName({'name-type': NAME_TYPE.MS_PRINCIPAL.value, 'name-string': cname})
		kdc_req_body_data['realm'] = 'WELLKNOWN:PKU2U'
		kdc_req_body_data['sname'] = PrincipalName({'name-type': NAME_TYPE.MS_PRINCIPAL.value, 'name-string': target})
		kdc_req_body_data['till']  = (now + datetime.timedelta(days=1)).replace(microsecond=0)
		kdc_req_body_data['rtime'] = (now + datetime.timedelta(days=1)).replace(microsecond=0)
		kdc_req_body_data['nonce'] = secrets.randbits(31)
		kdc_req_body_data['etype'] = [18,17] # 23 breaks...
		kdc_req_body_data['addresses'] = [HostAddress({'addr-type': 20, 'address': b'127.0.0.1'})] # not sure if this is needed
		kdc_req_body = KDC_REQ_BODY(kdc_req_body_data)


		checksum = hashlib.sha1(kdc_req_body.dump()).digest()
		
		authenticator = {}
		authenticator['cusec'] = now.microsecond
		authenticator['ctime'] = now.replace(microsecond=0)
		authenticator['nonce'] = secrets.randbits(31)
		authenticator['paChecksum'] = checksum
		

		dp = {}
		dp['p'] = self.diffie.p
		dp['g'] = self.diffie.g
		dp['q'] = 0 # mandatory parameter, but it is not needed

		pka = {}
		pka['algorithm'] = '1.2.840.10046.2.1'
		pka['parameters'] = keys.DomainParameters(dp)
		
		spki = {}
		spki['algorithm'] = keys.PublicKeyAlgorithm(pka)
		spki['public_key'] = self.diffie.get_public_key()

		
		authpack = {}
		authpack['pkAuthenticator'] = PKAuthenticator(authenticator)
		authpack['clientPublicValue'] = keys.PublicKeyInfo(spki)
		authpack['clientDHNonce'] = self.diffie.dh_nonce
		
		authpack = AuthPack(authpack)
		signed_authpack = self.sign_authpack(authpack.dump(), wrap_signed = False)
		
		# ??????? This is absolutely nonsense, 
		payload = length_encode(len(signed_authpack)) + signed_authpack
		payload = b'\x80' + payload
		signed_authpack = b'\x30' + length_encode(len(payload)) + payload
		
		pa_data_1 = {}
		pa_data_1['padata-type'] = PaDataType.PK_AS_REQ.value
		pa_data_1['padata-value'] = signed_authpack 

		asreq = {}
		asreq['pvno'] = 5
		asreq['msg-type'] = 10
		asreq['padata'] = [pa_data_1]
		asreq['req-body'] = kdc_req_body

		return AS_REQ(asreq).dump()	

	def build_apreq(self, asrep, session_key, cipher, subkey_data, krb_finished_data, flags = GSSAPIFlags.GSS_C_MUTUAL_FLAG | GSSAPIFlags.GSS_C_INTEG_FLAG  | GSSAPIFlags.GSS_C_EXTENDED_ERROR_FLAG):
		

		# TODO: https://www.ietf.org/rfc/rfc4757.txt
		#subkey_data = {}
		#subkey_data['keytype'] = Enctype.AES256
		#subkey_data['keyvalue'] = os.urandom(32)

		subkey_cipher = _enctype_table[subkey_data['keytype']]
		subkey_key = Key(subkey_cipher.enctype, subkey_data['keyvalue'])
		subkey_checksum = _checksum_table[16] # ChecksumTypes.hmac_sha1_96_aes256

		krb_finished_checksum_data = {}
		krb_finished_checksum_data['cksumtype'] = 16
		krb_finished_checksum_data['checksum'] = subkey_checksum.checksum(subkey_key, 41, krb_finished_data)

		krb_finished_data = {}
		krb_finished_data['gss-mic'] = Checksum(krb_finished_checksum_data)

		krb_finished = KRB_FINISHED(krb_finished_data).dump()

		a = 2
		extensions_data = a.to_bytes(4, byteorder='big', signed=True) + len(krb_finished).to_bytes(4, byteorder='big', signed=True) + krb_finished

		ac = AuthenticatorChecksum()
		ac.flags = flags
		ac.channel_binding = b'\x00'*16
		chksum = {}
		chksum['cksumtype'] = 0x8003
		chksum['checksum'] = ac.to_bytes() + extensions_data

		tii = LSAP_TOKEN_INFO_INTEGRITY()
		tii.Flags = 1
		tii.TokenIL = 0x00002000 # Medium integrity
		tii.MachineID = bytes.fromhex('7e303fffe6bff25146addca4fbddf1b94f1634178eb4528fb2731c669ca23cde')

		restriction_data = {}
		restriction_data['restriction-type'] = 0
		restriction_data['restriction'] = tii.to_bytes()
		restriction_data = KERB_AD_RESTRICTION_ENTRY(restriction_data)

		x = KERB_AD_RESTRICTION_ENTRYS([restriction_data]).dump()
		restrictions = AuthorizationData([{ 'ad-type' : 141, 'ad-data' : x}]).dump()

		

		now = datetime.datetime.now(datetime.timezone.utc)
		authenticator_data = {}
		authenticator_data['authenticator-vno'] = krb5_pvno 
		authenticator_data['crealm'] = Realm(asrep['crealm'])
		authenticator_data['cname'] = asrep['cname']
		authenticator_data['cusec'] = now.microsecond
		authenticator_data['ctime'] = now.replace(microsecond=0)
		authenticator_data['subkey'] = EncryptionKey(subkey_data)
		authenticator_data['seq-number'] = 682437742 #??? TODO: check this!
		authenticator_data['authorization-data'] = AuthorizationData([{'ad-type': 1, 'ad-data' : restrictions}])
		authenticator_data['cksum'] = Checksum(chksum)
		
		
		#print('Authenticator(authenticator_data).dump()')
		#print(Authenticator(authenticator_data).dump().hex())

		authenticator_data_enc = cipher.encrypt(session_key, 11, Authenticator(authenticator_data).dump(), None)
		
		ap_opts = ['mutual-required']

		ap_req = {}
		ap_req['pvno'] = krb5_pvno
		ap_req['msg-type'] = MESSAGE_TYPE.KRB_AP_REQ.value
		ap_req['ticket'] = Ticket(asrep['ticket'])
		ap_req['ap-options'] = APOptions(set(ap_opts))
		ap_req['authenticator'] = EncryptedData({'etype': session_key.enctype, 'cipher': authenticator_data_enc})
		
		#pprint('AP_REQ \r\n%s' % AP_REQ(ap_req).native)
		
		#print(AP_REQ(ap_req).dump().hex())
		#input()

		return AP_REQ(ap_req).dump()

	def sign_authpack(self, data, wrap_signed = False):
		if self.__hcert is not None:
			from minikerberos.common.windows.crypt32 import pkcs7_sign

			return pkcs7_sign(self.__hcert, data)
		return self.sign_authpack_native(self, data, wrap_signed)

	def sign_authpack_native(self, data, wrap_signed = False):
		"""
		Creating PKCS7 blob which contains the following things:

		1. 'data' blob which is an ASN1 encoded "AuthPack" structure
		2. the certificate used to sign the data blob
		3. the singed 'signed_attrs' structure (ASN1) which points to the "data" structure (in point 1)
		"""
		
		da = {}
		da['algorithm'] = algos.DigestAlgorithmId('1.3.14.3.2.26') # for sha1

		si = {}
		si['version'] = 'v1'
		si['sid'] = cms.IssuerAndSerialNumber({
			'issuer':  self.certificate.issuer,
			'serial_number':  self.certificate.serial_number,
		})


		si['digest_algorithm'] = algos.DigestAlgorithm(da)
		si['signed_attrs'] = [
			cms.CMSAttribute({'type': 'content_type', 'values': ['1.3.6.1.5.2.3.1']}), # indicates that the encap_content_info's authdata struct (marked with OID '1.3.6.1.5.2.3.1' is signed )
			cms.CMSAttribute({'type': 'message_digest', 'values': [hashlib.sha1(data).digest()]}), ### hash of the data, the data itself will not be signed, but this block of data will be.
		]
		si['signature_algorithm'] = algos.SignedDigestAlgorithm({'algorithm' : '1.2.840.113549.1.1.1'})
		si['signature'] = rsa_pkcs1v15_sign(self.privkey,  cms.CMSAttributes(si['signed_attrs']).dump(), "sha1")

		ec = {}
		ec['content_type'] = '1.3.6.1.5.2.3.1'
		ec['content'] = data

		sd = {}
		sd['version'] = 'v3'
		sd['digest_algorithms'] = [algos.DigestAlgorithm(da)] # must have only one
		sd['encap_content_info'] = cms.EncapsulatedContentInfo(ec)
		sd['certificates'] = [self.certificate]
		sd['signer_infos'] = cms.SignerInfos([cms.SignerInfo(si)])
		
		if wrap_signed is True:
			ci = {}
			ci['content_type'] = '1.2.840.113549.1.7.2' # signed data OID
			ci['content'] = cms.SignedData(sd)
			return cms.ContentInfo(ci).dump()

		return cms.SignedData(sd).dump()
		
	
	
	
	
	def decrypt_asrep(self, as_rep):
		

		
		def truncate_key(value, keysize):
			output = b''
			currentNum = 0
			while len(output) < keysize:
				currentDigest = hashlib.sha1(bytes([currentNum]) + value).digest()
				if len(output) + len(currentDigest) > keysize:
					output += currentDigest[:keysize - len(output)]
					break
				output += currentDigest
				currentNum += 1
			
			return output

		for pa in as_rep['padata']:
			if pa['padata-type'] == 17:
				pkasrep = PA_PK_AS_REP.load(pa['padata-value']).native
				break
		else:
			raise Exception('PA_PK_AS_REP not found!')

		sd = cms.SignedData.load(pkasrep['dhSignedData']).native
		keyinfo = sd['encap_content_info']
		if keyinfo['content_type'] != '1.3.6.1.5.2.3.2':
			raise Exception('Keyinfo content type unexpected value')
		authdata = KDCDHKeyInfo.load(keyinfo['content']).native
		pubkey = int(''.join(['1'] + [str(x) for x in authdata['subjectPublicKey']]), 2)		

		pubkey = int.from_bytes(core.BitString(authdata['subjectPublicKey']).dump()[7:], 'big', signed = False)
		shared_key = self.diffie.exchange(pubkey)
		
		server_nonce = pkasrep['serverDHNonce']
		fullKey = shared_key + self.diffie.dh_nonce + server_nonce

		etype = as_rep['enc-part']['etype']
		cipher = _enctype_table[etype]
		if etype == Enctype.AES256:
			t_key = truncate_key(fullKey, 32)
		elif etype == Enctype.AES128:
			t_key = truncate_key(fullKey, 16)
		elif etype == Enctype.RC4:
			raise NotImplementedError('RC4 key truncation documentation missing. it is different from AES')
			#t_key = truncate_key(fullKey, 16)
		

		key = Key(cipher.enctype, t_key)
		enc_data = as_rep['enc-part']['cipher']
		dec_data = cipher.decrypt(key, 3, enc_data)
		encasrep = EncASRepPart.load(dec_data).native
		cipher = _enctype_table[ int(encasrep['key']['keytype'])]
		session_key = Key(cipher.enctype, encasrep['key']['keyvalue'])
		return encasrep, session_key, cipher

	
	def get_metadata(self, target = None):
		


		if target is not None:
			if isinstance(target, str):
				target = [target]
		else:
			target = ['127.0.0.1']

		
		ci = {}
		ci['type'] = '2.5.4.3'
		ci['value'] = self.issuer

		a = Dunno1([ci])
		ci = Dunno2([a])

		info = {}
		info['pku2u'] = 'WELLKNOWN:PKU2U'
		info['clientInfo'] = PrincipalName({'name-type': NAME_TYPE.MS_PRINCIPAL.value, 'name-string': target})

		md = {}
		md['Info'] = Info(info)
		md['1'] = [CertIssuer({'data' : ci.dump()})]

		return MetaData(md).dump()