// SPDX-License-Identifier: GPL-3.0-only
// Copyright (c) 2022, Sylvain Huet, Ambermind
// Minimacy (r) System

use core.crypto.oid;;
use core.crypto.asn1;;
use core.crypto.rsa;;
use core.crypto.ec;;
use core.crypto.ed25519;;
use core.crypto.x25519;;
use core.crypto.hash;;
use core.crypto.des;;
use core.crypto.aes;;
use core.crypto.block;;
use core.crypto.pem;;
use core.util.base64;;

sum Key = rsaKey _, ecKey _, ed25519Key _, x25519Key _, aesKey _:Str, desKey _:Str, des3Key _:Str;;

const _DEBUG=false;;
fun rsaFromKey(key) = match key with rsaKey rsa -> rsa;;
fun ecFromKey(key) = match key with ecKey ec -> ec;;
fun ed25519FromKey(key) = match key with ed25519Key ed -> ed;;
fun x25519FromKey(key) = match key with x25519Key x -> x;;
fun aesFromKey(key) = match key with aesKey aes -> aes;;
fun desFromKey(key) = match key with desKey des -> des;;
fun des3FromKey(key) = match key with des3Key des3 -> des3;;

//-------------RSA-----------------------------
fun _rsaFromAsn1(asn1)=
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) -> content in
	if 9==listLength(content) then
	let content -> version::numbers in
	if 0==intFromAsn1(version) then
	let asn1ListOfBignum(numbers) -> modulus::pubExp::privExp::P::Q::dP::dQ::qInv::_ in
	rsaKey rsaFromPQ(pubExp, P, Q);;

fun _rsaPubFromAsn1(asn1)=
	if _DEBUG then asn1Echo("", asn1);
	let asn1ListOfBignum(listFromAsn1(asn1)) -> modulus::pubExp::_ in
	rsaKey rsaFromPublic(modulus, pubExp);;

fun _rsaPubFromSsh(pubExp, modulus)=
	rsaKey rsaFromPublic(bigDeserialize(modulus), bigDeserialize(pubExp));;

fun _forcePositive(number)= strConcat("\0", bigSerialize(number, nil));;		// start with \0 to force modulus to be considered as unsigned

fun asn1RsaKeyPub(modulus, pubExp)=
	asn1Seq
		(asn1Integer _forcePositive(modulus))::
		(asn1Integer bigSerialize(pubExp, nil))::nil;;

fun asn1RsaKeyPrivate(rsa)=
	asn1Seq
		(asn1Integer "\0")::
		(asn1Integer _forcePositive(rsa.modulusR))::
		(asn1Integer _forcePositive(rsa.pubExpR))::
		(asn1Integer _forcePositive(rsa.privExpR))::
		(asn1Integer _forcePositive(rsa.pR))::
		(asn1Integer _forcePositive(rsa.qR))::
		(asn1Integer _forcePositive(rsa.dPR))::
		(asn1Integer _forcePositive(rsa.dQR))::
		(asn1Integer _forcePositive(rsa.qInvR))::
		nil;;

fun asn1PrivateRsaEncryption(rsa)=
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier OID_rsaEncryption)::asn1Null::nil)::
		(asn1OctetString asn1Pack(asn1RsaKeyPrivate(rsa)))::nil;;

fun asn1PublicRsaEncryption(rsa)=
	let rsaModulus(rsa) -> modulus in
	let rsaPubExp(rsa) -> pubExp in
	asn1Seq
		(asn1Seq (asn1ObjectIdentifier OID_rsaEncryption)::asn1Null::nil)::
		(asn1BitString 0 asn1Pack(asn1RsaKeyPub(modulus, pubExp)))::nil;;

//-------------EC------------------------------
fun ecCurveFromOID(oid)=
	match oid with
		OID_secp256k1 -> ecSecp256k1(),
		OID_prime256v1-> ecSecp256r1(),
		OID_secp384r1 -> ecSecp384r1(),
		OID_secp521r1 -> ecSecp521r1();;

fun oidFromCurveName(name)=
	match name with
		"secp256k1" -> OID_secp256k1,
		"secp256r1" -> OID_prime256v1,
		"secp384r1" -> OID_secp384r1,
		"secp521r1" -> OID_secp521r1;;

fun _ecFromAsn1WithCurve(asn1, curve)=
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) -> content in
	if 3<=listLength(content) then
	let content ->  version::priv::pub::_ in
	if 1==intFromAsn1(version) then
	let ecCurveFromOID(curve) -> ec in
	ecKey ecKeyFromPrivate(ec, bigFromHex(hexFromStr(strFromAsn1(priv))));;

fun _ecFromAsn1(asn1)=
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) -> content in
	if 4<=listLength(content) then
	let content ->  version::priv::curve::pub::_ in
	if 1==intFromAsn1(version) then
	let ecCurveFromOID(strFromAsn1(head(listFromAsn1(curve)))) -> ec in
	ecKey ecKeyFromPrivate(ec, bigFromHex(hexFromStr(strFromAsn1(priv))));;

fun _ecPubFromAsn1(asn1, curve) =
	let ecCurveFromOID(strFromAsn1(curve)) -> ec in
	if ec<>nil then
	ecKey ecKeyFromPublic(ec, strFromAsn1(asn1));;

fun asn1EcKeyPrivate(ec)=
	asn1Seq
		(asn1Integer "\1")::
		(asn1OctetString bigSerialize(ecPrivate(ec), nil))::
		(asn1ClassConstructed ASN1_CONTEXT_SPECIFIC 1
			(asn1BitString 0 ecKeyPubStr(ec))::nil
		)::
		nil;;

fun asn1PrivateEcEncryption(ec)=
	let oidFromCurveName(ecKeyCurveName(ec)) -> curve in
	if curve<>nil then
	let ecKeyPubStr(ec) -> pubKey in
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier OID_ecPublicKey)::(asn1ObjectIdentifier curve)::nil)::
		(asn1OctetString asn1Pack(asn1EcKeyPrivate(ec)))::nil;;

fun asn1PublicEcEncryption(ec)=
	let oidFromCurveName(ecKeyCurveName(ec)) -> curve in
	if curve<>nil then
	let ecKeyPubStr(ec) -> pubKey in
	asn1Seq
		(asn1Seq (asn1ObjectIdentifier OID_ecPublicKey)::(asn1ObjectIdentifier curve)::nil)::
		(asn1BitString 0 pubKey)::nil;;

//-------------ed25519---------------------
fun _ed25519PubFromAsn1(raw)=
	ed25519Key ed25519FromPublic(raw);;

fun _ed25519PubFromSsh(raw)=
	ed25519Key ed25519FromPublic(raw);;

fun _ed25519FromAsn1(raw)=
	ed25519Key ed25519FromSecret(strFromAsn1(asn1Unpack(raw)));;
fun asn1PublicEd25519Encryption([secret, public])=
	asn1Seq
		(asn1Seq (asn1ObjectIdentifier OID_ed25519)::nil)::
		(asn1BitString 0 public)::nil;;

fun asn1PrivateEd25519Encryption([secret, public])=
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier OID_ed25519)::nil)::
		(asn1OctetString asn1Pack(asn1OctetString secret))::nil;;

//-------------x25519---------------------
fun _x25519PubFromAsn1(raw)=
	x25519Key x25519FromPublic(raw);;

fun _x25519FromAsn1(raw)=
	x25519Key x25519KeyPair(strFromAsn1(asn1Unpack(raw)));;

fun asn1PublicX25519Encryption([private, public])=
	asn1Seq
		(asn1Seq (asn1ObjectIdentifier OID_x25519)::nil)::
		(asn1BitString 0 public)::nil;;

fun asn1PrivateX25519Encryption([private, public])=
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier OID_x25519)::nil)::
		(asn1OctetString asn1Pack(asn1OctetString private))::nil;;

//-------------AES-----------------------------
// not standard, modeled from ED25519 with a credible OID
fun asn1PrivateAesEncryption(private)=
	let match strLength(private) with
		16 -> OID_AES128_CBC,
		24 -> OID_AES192_CBC,
		32 -> OID_AES256_CBC
	-> oid in
	if oid<>nil then
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier oid)::nil)::
		(asn1OctetString asn1Pack(asn1OctetString private))::nil;;
fun aesFromAsn1(raw)= aesKey strFromAsn1(asn1Unpack(raw));;

//-------------DES-----------------------------
// not standard, modeled from ED25519 with a credible OID
fun asn1PrivateDesEncryption(private)=
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier OID_DES_CBC)::nil)::
		(asn1OctetString asn1Pack(asn1OctetString private))::nil;;
fun desFromAsn1(raw)= desKey strFromAsn1(asn1Unpack(raw));;

//-------------DES3-----------------------------
// not standard, modeled from ED25519 with a credible OID
fun asn1PrivateDes3Encryption(private)=
	asn1Seq
		(asn1Integer "\0")::
		(asn1Seq (asn1ObjectIdentifier OID_DES_3CBC)::nil)::
		(asn1OctetString asn1Pack(asn1OctetString private))::nil;;
fun des3FromAsn1(raw)= des3Key strFromAsn1(asn1Unpack(raw));;

//-------------encrypted private key---------------------
fun _keyDerivation(pwd, keyLen, asn1) =
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) -> algo::params::_ in
	match strFromAsn1(algo) with
		OID_PBKDF2 -> let listFromAsn1(params) ->salt::count::hashFunction::_ in
			let listFromAsn1(hashFunction) -> algo::params in
			match strFromAsn1(algo) with
				OID_hmacWithSHA256 -> pbkdf2(pwd, strFromAsn1(salt), intFromAsn1(count), keyLen, #hmacSha256);;

fun _cypherExtract(asn1) =
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) -> algo::params in
	match strFromAsn1(algo) with
		OID_AES256_CBC  -> [lambda(key, bin) = aesDecryptCbc( key, strFromAsn1(head(params)), bin), 32],
		OID_des_ede3_cbc-> [lambda(key, bin) = des3DecryptCbc(key, strFromAsn1(head(params)), bin), 24];;

fun _pBES2KeyFromAsn1(pwd, derivation::cipher::_, raw)=
	let _cypherExtract(cipher) -> [fCipher, keyLen] in
	let _keyDerivation(pwd, keyLen, derivation) -> key in
	let call fCipher(key, strFromAsn1(raw)) -> data in
	strLeft(data, unPaddingPKCS_5(data));;

fun asn1Encrypt(pwd, data)=	// with PBKDF2, hmacWithSHA256, AES256-CBC
	let strRand(8) -> saltPbkdf2 in
	let 2048 -> count in
	let pbkdf2(pwd, saltPbkdf2, count, 32, #hmacSha256) -> key in
	let strRand(16) -> iv in
	let strConcat(data, paddingPKCS_5(data, AES_BLOCK)) -> data in
	let aesEncryptCbc(key, iv, data) -> data in
	asn1Seq
		(asn1Seq 
			(asn1ObjectIdentifier OID_PBES2)::
			(asn1Seq 
				(asn1Seq 
					(asn1ObjectIdentifier OID_PBKDF2)::
					(asn1Seq 
						(asn1OctetString saltPbkdf2)::
						(asn1Integer strInt16Msb(count))::
						(asn1Seq 
							(asn1ObjectIdentifier OID_hmacWithSHA256)::
							asn1Null::
							nil)::
						nil)::
					nil)::
				(asn1Seq 
					(asn1ObjectIdentifier OID_AES256_CBC)::
					(asn1OctetString iv)::
					nil)::
				nil)::
			nil)::
		(asn1OctetString data)::nil;;

fun asn1Decrypt(asn1, pwd)=
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) ->(encryption::raw::_) in
	let listFromAsn1(encryption) -> algo::next::_ in
	match strFromAsn1(algo) with
		OID_PBES2 -> _pBES2KeyFromAsn1(pwd, listFromAsn1(next), raw);;

//-------------ssh key extraction
fun _parseSSH(data, i)=
	if i<strLength(data) then
	let strRead32Msb(data, i) -> len in
	if len>=0 then strSlice(data, i+4, len)::_parseSSH(data, i+4+len);;
	
fun keyFromPublicSSH(b64) =
	let _parseSSH(b64Decode(b64), 0) -> type::a::b::_ in
	match type with
	"ssh-rsa" -> _rsaPubFromSsh(a, b),
	"ssh-ed25519" -> _ed25519PubFromSsh(a);;

//-------------generic key extraction
fun keyFromAsn1(asn1, pwd)=
	if _DEBUG then asn1Echo("", asn1);
	let listFromAsn1(asn1) -> asn1List in
	match head(asn1List) with	// we try to guess the kind of key, based on the asn1 structure
		asn1Integer _ -> (// private key
			let asn1List ->(version::encryption::raw::_) in
			if 0==intFromAsn1(version) then
			let listFromAsn1(encryption) -> encryptionMode::curve::_ in
			match strFromAsn1(encryptionMode) with
				OID_rsaEncryption -> _rsaFromAsn1(asn1Unpack(strFromAsn1(raw))),
				OID_ecPublicKey -> _ecFromAsn1WithCurve(asn1Unpack(strFromAsn1(raw)), strFromAsn1(curve)),
				OID_ed25519 -> _ed25519FromAsn1(strFromAsn1(raw)),
				OID_x25519 -> _x25519FromAsn1(strFromAsn1(raw)),
				OID_DES_CBC -> desFromAsn1(strFromAsn1(raw)),	// non standard
				OID_DES_3CBC -> des3FromAsn1(strFromAsn1(raw)),	// non standard
				OID_AES128_CBC -> aesFromAsn1(strFromAsn1(raw)),	// non standard
				OID_AES192_CBC -> aesFromAsn1(strFromAsn1(raw)),	// non standard
				OID_AES256_CBC -> aesFromAsn1(strFromAsn1(raw))	// non standard
		),
		_ -> match head(tail(asn1List)) with
			asn1BitString _ _-> (	// public key
				let asn1List ->(encryption::raw::_) in
				let listFromAsn1(encryption) -> lEncryption in
				let strFromAsn1(head(lEncryption)) -> encryptionAlgo in 
				match encryptionAlgo with
					OID_rsaEncryption -> _rsaPubFromAsn1(asn1Unpack(strFromAsn1(raw))),
					OID_ecPublicKey -> _ecPubFromAsn1(raw, head(tail(lEncryption))),
					OID_ed25519 -> _ed25519PubFromAsn1(strFromAsn1(raw)),
					OID_x25519 -> _x25519PubFromAsn1(strFromAsn1(raw))
			),
			asn1OctetString _ -> (	// encrypted private key
				let asn1Decrypt(asn1, pwd) -> der in
				keyFromDER(der, nil)
			),
			asn1Seq _ -> 	// certificate or csr
				let asn1List ->(infoBlock::_) in
				let listFromAsn1(infoBlock) -> infoBlockContent in
				let if listLength(infoBlockContent)<7 then 2 else 6 -> index in	// Hack! 2:csr, 6:certificate
				let listGet(infoBlockContent, index) -> publicKey in
				keyFromAsn1(publicKey, pwd)
			;;

fun derDecrypt(der, pwd) =
	let asn1Unpack(der) -> asn1 in
	asn1Decrypt(asn1, pwd);;

fun keyFromDER(der, pwd) =
	let asn1Unpack(der) -> asn1 in
	keyFromAsn1(asn1, pwd);;

// ssh
fun sshParseInts(data, i, n)=
	if n>0 && i<strLength(data) then
	let strRead32Msb(data, i) -> val in
	val::sshParseInts(data, i+4, n-1);;

fun sshParseVals(data, i, n)=
	if n>0 && i<strLength(data) then
	let strRead32Msb(data, i) -> len in
	strSlice(data, i+4, len)::sshParseVals(data, i+4+len, n-1);;

fun sshParseValsNext(data, i, n)=	// returns the offset after the last parsed value
	if n<=0 || i>strLength(data) then i
	else 
	let strRead32Msb(data, i) -> len in
	sshParseValsNext(data, i+4+len, n-1);;

fun _keyFromOpenSSH(data)=	// https://dnaeon.github.io/openssh-private-key-binary-format/
	let strPos(data, "\0", 0) -> i in
	if i<>nil then
	let strLeft(data, i) -> magic in
	if magic=="openssh-key-v1" then	// the only recognized value for now
	let sshParseVals(data, i+1, 3) -> ciphername::kdfname::kdfoptions::_ in
	if ciphername=="none" && kdfname=="none" then
	let sshParseValsNext(data, i+1, 3) -> i in
	let i+4 -> i in	// we skip the msb32int providing the number of keys, we need only the first one
	let sshParseVals(data, i, 2) -> pubBlob::privBlob::_ in
	let sshParseVals(privBlob, 2*4, 3) -> type::pub::publicThenPrivate::_ in
	if type=="ssh-ed25519" then	// the only recognized value for now
	ed25519Key ed25519FromSecret(strLeft(publicThenPrivate, 32));;

// PEM
fun keyFromPEM(pem, pwd)=
	let head(pemDecode(pem, pwd)) ->[name, der] in
	if name=="OPENSSH PRIVATE KEY" then _keyFromOpenSSH(der)
	else
	let asn1Unpack(der) -> asn1 in
	match name with
		"PRIVATE KEY"->keyFromAsn1(asn1, pwd),
		"ENCRYPTED PRIVATE KEY"->keyFromAsn1(asn1, pwd),
		"RSA PRIVATE KEY"->_rsaFromAsn1(asn1),
		"EC PRIVATE KEY"->_ecFromAsn1(asn1),
		"PUBLIC KEY"->keyFromAsn1(asn1, pwd),
		"CERTIFICATE"->keyFromAsn1(asn1, pwd),
		"CERTIFICATE REQUEST"->keyFromAsn1(asn1, pwd);;

fun rsaFromPEM(pem, pwd)= rsaFromKey(keyFromPEM(pem, pwd));;
fun ecKeyFromPEM(pem, pwd)= ecFromKey(keyFromPEM(pem, pwd));;

fun keyIsPrivate(key)= match key with
	ecKey ec     -> ecIsPrivate(ec),
	rsaKey rsa   -> rsaIsPrivate(rsa),
	ed25519Key ed-> ed25519KeyIsPrivate(ed),
	x25519Key x  -> x25519KeyIsPrivate(x),
	aesKey aes   -> true,
	desKey des   -> true,
	des3Key des  -> true;;

fun asn1PrivateEncryption(key)= match key with
	ecKey ec     -> asn1PrivateEcEncryption(ec),
	rsaKey rsa   -> asn1PrivateRsaEncryption(rsa),
	ed25519Key ed-> asn1PrivateEd25519Encryption(ed),
	x25519Key x  -> asn1PrivateX25519Encryption(x),
	aesKey aes   -> asn1PrivateAesEncryption(aes),
	desKey des   -> asn1PrivateDesEncryption(des),
	des3Key des3  -> asn1PrivateDes3Encryption(des3);;
fun derFromPrivateKey(key) = asn1Pack(asn1PrivateEncryption(key));;
fun pemFromPrivateKey(key, pwd) = 
	let derFromPrivateKey(key) -> data in
	if pwd==nil then pemMake("PRIVATE KEY", data)
	else pemMake("ENCRYPTED PRIVATE KEY", asn1Pack(asn1Encrypt(pwd, data)));;

fun asn1PublicEncryption(key)= match key with
	ecKey ec     -> asn1PublicEcEncryption(ec),
	rsaKey rsa   -> asn1PublicRsaEncryption(rsa),
	ed25519Key ed-> asn1PublicEd25519Encryption(ed),
	x25519Key  x -> asn1PublicX25519Encryption(x);;
fun derFromPublicKey(key) = asn1Pack(asn1PublicEncryption(key));;
fun pemFromPublicKey(key) = pemMake("PUBLIC KEY", derFromPublicKey(key));;

fun publicKeyHash(key) = sha256(derFromPublicKey(key));;

fun keyDump(key)= match key with
	ecKey ec  -> (ecKeyDump(ec); nil),
	rsaKey rsa-> (rsaDump(rsa); nil);
	key;;