-- SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
-- SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later

--[==[
Cryptographic primitives and utilities built on LITLS. Provides
encoding/decoding (hex, base64, base64url), hashing (SHA-256, HMAC),
ECC P-256 and Ed25519 key generation, signing, and verification,
as well as X.509 certificate parsing and CSR generation.
]==]

local core = require("crypto.core")
local litls = require("litls.core")
local std = require("std")
local json = require("cjson.safe")

---! Convert a binary string to its hexadecimal representation
---@ bin_to_hex(`bin`{.str}) -> `hex`{.str}
local bin_to_hex = function(bin)
	local hex = string.gsub(bin, ".", function(c)
		return string.format("%02x", string.byte(c))
	end)
	return hex
end

---! Convert a hexadecimal string to binary
---@ hex_to_bin(`hex`{.str}) -> `bin`{.str}
local hex_to_bin = function(hex)
	local bin = string.gsub(hex, "..", function(h)
		return string.char(tonumber(h, 16))
	end)
	return bin
end

---! Compute SHA-1 hash of data
---@ sha1(`data`{.str}) -> `hash`{.str}
local sha1 = function(data)
	return core.sha1(data)
end

---! Compute SHA-256 hash of data
---@ sha256(`data`{.str}) -> `hash`{.str}
local sha256 = function(data)
	return core.sha256(data)
end

---! Compute HMAC-SHA256
---@ hmac(`secret`{.str}, `msg`{.str}) -> `mac`{.str}
local hmac = function(secret, msg)
	return core.hmac(secret, msg)
end

---! Encode a string to base64
---@ b64_encode(`str`{.str}) -> `encoded`{.str}
local b64_encode = function(str)
	return core.base64_encode(str)
end

---! Decode a base64-encoded string
---@ b64_decode(`str`{.str}) -> `decoded`{.str}
local b64_decode = function(str)
	return core.base64_decode(str)
end

---! Encode a string to base64url (URL-safe, no padding)
---@ b64url_encode(`str`{.str}) -> `encoded`{.str}
local b64url_encode = function(str)
	local b64_str = core.base64_encode(str)
	return b64_str:gsub("[+/=]", { ["+"] = "-", ["/"] = "_", ["="] = "" })
end

---! Decode a base64url-encoded string
---@ b64url_decode(`str`{.str}) -> `decoded`{.str}
local b64url_decode = function(str)
	str = str .. string.rep("=", (4 - (#str % 4)) % 4)
	return b64_decode(str:gsub("[-_]", { ["-"] = "+", ["_"] = "/" }))
end

---! JSON-encode a table and base64url-encode the result
---@ b64url_encode_json(`tbl`{.tbl}) -> `encoded`{.str}
local b64url_encode_json = function(tbl)
	local str = json.encode(tbl)
	return b64url_encode(str)
end

---! Generate an ECC P-256 key pair
---@ ecc_generate_key() -> `key_obj`{.tbl}
--[===[
Returns a key object table with raw binary string fields:
`private` (scalar d), `public` (uncompressed point),
`x` (32-byte affine x), `y` (32-byte affine y).
]===]
local ecc_generate_key = function()
	local key, pub, x, y = core.ecc_generate_key()
	local key_obj = {
		private = key,
		public = pub,
		x = x,
		y = y,
	}
	return key_obj
end

---! Save an ECC key object to a file in JWK format
---@ ecc_save_key(`key_obj`{.tbl}, `key_file`{.str}) -> `ok`{.bool .or_nil}, `err`{.str .err}
--[===[
Serializes the key object to JSON with base64url-encoded fields
(`x`, `y`, `d`, `pub`) and an optional `kid` field, then writes
the result to the given file path.
]===]
local ecc_save_key = function(key_obj, key_file)
	if not key_file then
		return nil, "no filename provided"
	end
	local jwk, err = json.encode({
		x = b64url_encode(key_obj.x),
		y = b64url_encode(key_obj.y),
		d = b64url_encode(key_obj.private),
		pub = b64url_encode(key_obj.public),
		kid = key_obj.kid,
	})
	if not jwk then
		return nil, "failed to encode key: " .. err
	end
	return std.fs.write_file(key_file, jwk)
end

---! Load an ECC key object from a JWK file
---@ ecc_load_key(`key_file`{.str}) -> `key_obj`{.tbl .or_nil}, `err`{.str .err}
--[===[
Reads a JWK JSON file written by `ecc_save_key` and returns a key
object table with raw binary fields: `private`, `public`, `x`, `y`,
and optionally `kid`.
]===]
local ecc_load_key = function(key_file)
	local content, err = std.fs.read_file(key_file)
	if not content then
		return nil, err
	end
	local jwk, err = json.decode(content)
	if not jwk then
		return nil, "failed to decode the key: " .. err
	end
	local key_obj = {
		private = b64url_decode(jwk.d),
		public = b64url_decode(jwk.pub),
		x = b64url_decode(jwk.x),
		y = b64url_decode(jwk.y),
		kid = jwk.kid,
	}
	return key_obj
end

---! Sign a message using an ECC P-256 private key
---@ ecc_sign(`key`{.str}, `pub_key`{.str}, `msg`{.str}) -> `sig`{.str .or_nil}, `err`{.str .err}
--[===[
All arguments are raw binary strings. `key` is the private scalar,
`pub_key` is the uncompressed public point. Returns a 64-byte
signature in r||s format (32 bytes each, big-endian).
]===]
local ecc_sign = function(key, pub_key, msg)
	return core.ecc_sign(key, pub_key, msg)
end

---! Verify an ECC P-256 signature
---@ ecc_verify(`pub_key`{.str}, `msg`{.str}, `sig`{.str}) -> `ok`{.bool .or_nil}, `err`{.str .err}
--[===[
Counterpart to `ecc_sign`. `pub_key` is the raw uncompressed public
point, `sig` is the 64-byte r||s signature. Returns true on success,
nil and an error message on failure.
]===]
local ecc_verify = function(pub_key, msg, sig)
	return core.ecc_verify(pub_key, msg, sig)
end

---! Generate an Ed25519 key pair
---@ ed25519_generate_key() -> `private_key`{.str}, `public_key`{.str}
local ed25519_generate_key = function()
	return core.ed25519_generate_key()
end

---! Sign a message using an Ed25519 private key
---@ ed25519_sign(`key`{.str}, `msg`{.str}) -> `sig`{.str .or_nil}, `err`{.str .err}
local ed25519_sign = function(key, msg)
	return core.ed25519_sign(key, msg)
end

---! Verify an Ed25519 signature
---@ ed25519_verify(`pub_key`{.str}, `msg`{.str}, `sig`{.str}) -> `ok`{.bool .or_nil}, `err`{.str .err}
local ed25519_verify = function(pub_key, msg, sig)
	return core.ed25519_verify(pub_key, msg, sig)
end

---! Generate a PKCS#10 certificate signing request
---@ generate_csr(`key`{.str}, `pub_key`{.str}, `domain`{.str}, `alt_names`{.tbl .opt}) -> `csr`{.str .or_nil}, `err`{.str .err}
--[===[
Creates a DER-encoded CSR for the given domain using raw binary ECC
key material. `alt_names`, when provided, is an array of additional
DNS names to include as Subject Alternative Names.
]===]
local generate_csr = function(key, pub_key, domain, alt_names)
	if not key or not pub_key then
		return nil, "you must provide a key"
	end
	if not domain then
		return nil, "domain name is required"
	end
	if alt_names and type(alt_names) ~= "table" then
		return nil, "alternative names must be a table"
	end
	return core.generate_csr(key, pub_key, domain, alt_names)
end

---! Convert an ECC key object from DER to PEM format
---@ der_to_pem_ecc_key(`key_obj`{.tbl}) -> `pem`{.str .or_nil}, `err`{.str .err}
--[===[
Accepts a key object table with `private` and `public` raw binary
fields (as returned by `ecc_generate_key`) and returns a
PEM-encoded EC private key string.
]===]
local der_to_pem_ecc_key = function(key_obj)
	if not key_obj or not key_obj.private or not key_obj.public then
		return nil, "invalid key object"
	end
	return core.der_to_pem_ecc_key(key_obj.private, key_obj.public)
end

---! Parse an X.509 certificate
---@ parse_x509_cert(`cert`{.str}) -> `cert_info`{.tbl .or_nil}, `err`{.str .err}
--[===[
Accepts a certificate in PEM or DER format. PEM armor is stripped
automatically. Returns a table with certificate fields such as
subject, issuer, serial, validity dates, and extensions.
]===]
local parse_x509_cert = function(cert)
	if not cert then
		return nil, "certificate is required"
	end
	local pem_start = std.escape_magic_chars("-----BEGIN CERTIFICATE-----")
	local pem_end = std.escape_magic_chars("-----END CERTIFICATE-----")
	local cert_der = cert
	if cert:match("^" .. pem_start) then
		local cert_b64 = cert:match("^" .. pem_start .. "(.-)" .. pem_end)
		cert_b64 = cert_b64:gsub("[\r\n%s]", "")
		cert_der = b64_decode(cert_b64)
	end
	local cert_info = core.parse_x509_cert(cert_der)
	if not cert_info then
		return nil, "failed to parse certificate"
	end
	return cert_info
end

-- ── SSH key helpers ──────────────────────────────────────────────────

---! Parse an ssh-ed25519 public key string
---@ ssh_parse_pubkey(`content`{.str}) -> `pubkey`{.str .or_nil}, `err`{.str .err}
--[===[
Parses an `ssh-ed25519` public key string (the format found in `.pub`
files or authorized_keys) and extracts the raw 32-byte Ed25519 public
key. Returns nil and an error if the content is not a valid
ssh-ed25519 key.
]===]
local ssh_parse_pubkey = function(content)
	local b64 = content:match("^ssh%-ed25519%s+(%S+)")
	if not b64 then
		return nil, "not an ssh-ed25519 public key"
	end
	local blob = core.base64_decode(b64)
	if not blob or #blob < 4 then
		return nil, "invalid key blob"
	end
	local pos = 1
	local type_len = blob:byte(pos) * 16777216
		+ blob:byte(pos + 1) * 65536
		+ blob:byte(pos + 2) * 256
		+ blob:byte(pos + 3)
	pos = pos + 4
	if #blob < pos + type_len - 1 then
		return nil, "truncated key type"
	end
	local key_type = blob:sub(pos, pos + type_len - 1)
	pos = pos + type_len
	if key_type ~= "ssh-ed25519" then
		return nil, "not an ed25519 key"
	end
	if #blob < pos + 3 then
		return nil, "truncated public key length"
	end
	local pk_len = blob:byte(pos) * 16777216
		+ blob:byte(pos + 1) * 65536
		+ blob:byte(pos + 2) * 256
		+ blob:byte(pos + 3)
	pos = pos + 4
	if pk_len ~= 32 then
		return nil, "invalid public key length: " .. tostring(pk_len)
	end
	if #blob < pos + 31 then
		return nil, "truncated public key"
	end
	return blob:sub(pos, pos + 31)
end

---! Parse an OpenSSH Ed25519 private key from PEM content
---@ ssh_parse_privkey(`pem`{.str}) -> `seed`{.str .or_nil}, `pubkey`{.str .or_nil .err}
--[===[
Parses an unencrypted OpenSSH Ed25519 private key (PEM format with
`-----BEGIN OPENSSH PRIVATE KEY-----` armor). Returns the 32-byte
seed and 32-byte public key, or nil and an error message.
]===]
local ssh_parse_privkey = function(pem)
	return litls.openssh_parse_ed25519(pem)
end

---! Load an Ed25519 public key from an OpenSSH .pub file
---@ ssh_load_pubkey(`path`{.str}) -> `pubkey`{.str .or_nil}, `err`{.str .err}
local ssh_load_pubkey = function(path)
	local content, err = std.fs.read_file(path)
	if not content then
		return nil, "cannot open public key file: " .. (err or path)
	end
	local line = content:match("^[^\n]+")
	if not line then
		return nil, "empty public key file"
	end
	return ssh_parse_pubkey(line)
end

---! Load an Ed25519 key from an OpenSSH private key file
---@ ssh_load_key(`path`{.str}) -> `seed`{.str .or_nil}, `pubkey`{.str .or_nil .err}
local ssh_load_key = function(path)
	local content, err = std.fs.read_file(path)
	if not content then
		return nil, "cannot open key file: " .. (err or path)
	end
	return ssh_parse_privkey(content)
end

---! Compute the hex-encoded SHA-256 fingerprint of an Ed25519 public key
---@ ssh_pubkey_fingerprint(`pubkey`{.str}) -> `fingerprint`{.str}
local ssh_pubkey_fingerprint = function(pubkey)
	return bin_to_hex(core.sha256(pubkey))
end

return {
	bin_to_hex = bin_to_hex,
	hex_to_bin = hex_to_bin,
	b64_encode = b64_encode,
	b64_decode = b64_decode,
	b64url_encode = b64url_encode,
	b64url_decode = b64url_decode,
	b64url_encode_json = b64url_encode_json,
	sha1 = sha1,
	sha256 = sha256,
	hmac = hmac,
	ecc_generate_key = ecc_generate_key,
	ecc_load_key = ecc_load_key,
	ecc_save_key = ecc_save_key,
	ecc_sign = ecc_sign,
	ecc_verify = ecc_verify,
	ed25519_generate_key = ed25519_generate_key,
	ed25519_sign = ed25519_sign,
	ed25519_verify = ed25519_verify,
	generate_csr = generate_csr,
	der_to_pem_ecc_key = der_to_pem_ecc_key,
	parse_x509_cert = parse_x509_cert,
	ssh_parse_pubkey = ssh_parse_pubkey,
	ssh_parse_privkey = ssh_parse_privkey,
	ssh_load_pubkey = ssh_load_pubkey,
	ssh_load_key = ssh_load_key,
	ssh_pubkey_fingerprint = ssh_pubkey_fingerprint,
}