import {crypto} from "jsr:@std/crypto";
import {decodeBase64} from "jsr:@std/encoding/base64";
import {encodeBase64Url} from "jsr:@std/encoding/base64url";
import {decodeHex, encodeHex} from "jsr:@std/encoding/hex";

// IMPORTANT: ONLY supports ECDSA P256 and P384

// Reference:
// * https://github.com/openssh/openssh-portable/blob/master/PROTOCOL.sshsig
// * https://docs.deno.com/examples/hex_base64_encoding/
// * https://jsr.io/@std/encoding/1.0.6/base64url.ts
// * https://docs.deno.com/examples/hashing/
// * https://developer.mozilla.org/en-US/docs/Web/API/SubtleCrypto/importKey
// * https://developer.mozilla.org/en-US/docs/Web/API/SubtleCrypto/verify

// #define MAGIC_PREAMBLE "SSHSIG"
// #define SIG_VERSION    0x01
//
//         byte[6]   MAGIC_PREAMBLE
//         uint32    SIG_VERSION
//         string    publickey                   // @see public key format
//         string    namespace
//         string    reserved                    // always empty
//         string    hash_algorithm
//         string    signature                   // @see signature format
//
// public key format
//         string    ecdsa-sha2-nistp256|384|521
//         string    nistp256|384|521
//         string    ec public key point
//
// signature format
//         string    ecdsa-sha2-nistp256|384|521
//         string    signature of ssh sign data  // @see ec signature format
//
// ec signature format
//         string    r
//         string    s
//
// ssh sign data
//         byte[6]   MAGIC_PREAMBLE
//         string    namespace
//         string    reserved                    // always empty
//         string    hash_algorithm
//         string    input digest hash with @hash_algorithm

const ENCODING_UTF8 = "UTF-8";

class BinaryWriter {
    buffers: Array<Uint8Array>;

    constructor() {
        this.buffers = [];
    }

    toArray(): Uint8Array {
        let totalLen = 0;
        for (let i = 0; i < this.buffers.length; i++) {
            totalLen += this.buffers[i].byteLength;
        }
        let merged = new Uint8Array(totalLen);
        let offset = 0;
        for (let i = 0; i < this.buffers.length; i++) {
            merged.set(this.buffers[i], offset);
            offset += this.buffers[i].byteLength;
        }
        return merged;
    }

    writeString(byte: Uint8Array) {
        this.writeUint32(byte.byteLength);
        this.writeBytes(byte);
    }

    writeStringFromString(str: string) {
        const bytes = new TextEncoder().encode(str);
        this.writeString(bytes);
    }

    writeLengthFromString(str: string) {
        const bytes = new TextEncoder().encode(str);
        this.writeBytes(bytes);
    }

    writeUint32(num: number) {
        let dataView = new DataView(new ArrayBuffer(4), 0);
        dataView.setUint32(0, num);
        this.writeBytes(new Uint8Array(dataView.buffer));
    }

    writeNumber(num: number) {
        let n = new Uint8Array(1)
        n[0] = num;
        this.writeBytes(n);
    }

    writeBytes(bytes: Uint8Array) {
        this.buffers.push(bytes);
    }
}

class BinaryReader {
    buffer: Uint8Array;
    pos: number;

    constructor(buffer: Uint8Array) {
        this.buffer = buffer;
        this.pos = 0;
    }

    readStringToString(): string {
        const buff = this.readString();
        return new TextDecoder(ENCODING_UTF8).decode(buff);
    }

    readLengthToString(len: number): string {
        const buff = this.readBytes(len);
        return new TextDecoder(ENCODING_UTF8).decode(buff);
    }

    readString(): Uint8Array {
        const len = this.readUint32();
        return this.readBytes(len);
    }

    readUint32(): number {
        const num = this.readBytes(4);
        const dataView = new DataView(num.buffer, 0);
        return dataView.getUint32(0);
    }

    readBytes(len: number): Uint8Array {
        if (this.buffer.byteLength < (this.pos + len)) {
            throw `Bad length ${this.buffer.byteLength} < ${this.pos} + ${len}`;
        }
        const bytes = this.buffer.slice(this.pos, this.pos + len);
        this.pos += len;
        return bytes;
    }
}

class SshSignature {
    publicKey: SshPublicKey;
    namespace: string;
    hashAlgorithm: string;
    signature: SshSignatureValue;

    constructor(publicKey: SshPublicKey, namespace: string, hashAlgorithm: string, signature: SshSignatureValue) {
        this.publicKey = publicKey;
        this.namespace = namespace;
        this.hashAlgorithm = hashAlgorithm;
        this.signature = signature;
    }

    static async parseFile(filename: string): Promise<SshSignature> {
        const pem = await Deno.readTextFile(filename);
        return SshSignature.parsePem(pem);
    }

    static parsePem(pem: string): SshSignature {
        const buffer = parsePemToArray(pem);
        return SshSignature.parse(buffer);
    }

    static parse(buffer: Uint8Array): SshSignature {
        let reader = new BinaryReader(buffer);
        const sshSig = reader.readLengthToString(6);
        if (sshSig !== "SSHSIG") {
            throw `Bad SSH signature magic, expect: SSHSIG, actual: ${sshSig}`;
        }
        let sshVer = reader.readUint32();
        if (sshVer !== 1) {
            throw `Bad SSH signature version, expect: 1, actual: ${sshVer}`;
        }
        const publicKeyBytes = reader.readString();
        const publicKey = SshPublicKey.parse(publicKeyBytes);
        const namespace = reader.readStringToString();
        const reserved = reader.readStringToString();
        if (reserved !== "") {
            throw `Bad SSH signature reserved value, expect empty, actual: ${reserved}`;
        }
        const hashAlgorithm = reader.readStringToString();
        const signatureBytes = reader.readString();
        const signature = SshSignatureValue.parse(signatureBytes);
        return new SshSignature(publicKey, namespace, hashAlgorithm, signature);
    }

    async verifyString(data: string): Promise<boolean> {
        const digest = await this._digestString(data);
        return await this.verifyDigest(digest);
    }

    async verifyFile(filename: string): Promise<boolean> {
        const digest = await this._digestFile(filename);
        return await this.verifyDigest(digest);
    }

    async verifyDigest(digest: Uint8Array): Promise<boolean> {
        const publicKey = await this.publicKey.importJwk();
        const signature = this.signature.toRs();
        const signatureData = this.calculateSignatureData(digest);
        return await crypto.subtle.verify(
            {
                name: "ECDSA",
                hash: {name: (this.publicKey.algorithm === "nistp256") ? "SHA-256" : "SHA-384"},
            },
            publicKey,
            signature,
            signatureData,
        );
    }

    calculateSignatureData(digest: Uint8Array): Uint8Array {
        if (this.getHashAlgorithmLength() !== digest.byteLength) {
            throw `Bad digest length, expect: ${this.getHashAlgorithmLength()}, acutal: ${digest.byteLength}`;
        }
        let writer = new BinaryWriter();
        writer.writeLengthFromString("SSHSIG");
        writer.writeStringFromString(this.namespace);
        writer.writeStringFromString("");
        writer.writeStringFromString(this.hashAlgorithm);
        writer.writeString(digest);
        return writer.toArray();
    }

    getHashAlgorithmLength(): number {
        switch (this.hashAlgorithm.toLowerCase()) {
            case "sha512":
            case "sha-512":
                return 512 / 8;
            case "sha384":
            case "sha-384":
                return 384 / 8;
            case "sha256":
            case "sha-256":
                return 256 / 8;
            default:
                throw `Unknown hash algorithm: ${this.hashAlgorithm}`;
        }
    }

    async _digestString(data: string): Promise<Uint8Array> {
        return await digestString(data, this.hashAlgorithm);
    }

    async _digestFile(filename: string): Promise<Uint8Array> {
        return await digestFile(filename, this.hashAlgorithm);
    }
}

class SshSignatureValue {
    signatureAlgorithm: string;
    ecSignatureR: Uint8Array;
    ecSignatureS: Uint8Array;

    constructor(signatureAlgorithm: string, ecSignatureR: Uint8Array, ecSignatureS: Uint8Array) {
        this.signatureAlgorithm = signatureAlgorithm;
        this.ecSignatureR = ecSignatureR;
        this.ecSignatureS = ecSignatureS;
    }

    static parse(buffer: Uint8Array): SshSignatureValue {
        let reader = new BinaryReader(buffer);
        const signatureAlgorithm = reader.readStringToString();
        const signatureEc = reader.readString();
        let signatureEcReader = new BinaryReader(signatureEc);
        const r = signatureEcReader.readString();
        const s = signatureEcReader.readString();
        return new SshSignatureValue(signatureAlgorithm, r, s);
    }

    // crypto.subtle.sign's signature is R and S
    toRs(): Uint8Array {
        let writer = new BinaryWriter();
        if (this.ecSignatureR.byteLength === 0x21) {
            writer.writeBytes(this.ecSignatureR.slice(1));
        } else {
            writer.writeBytes(this.ecSignatureR);
        }
        if (this.ecSignatureS.byteLength === 0x21) {
            writer.writeBytes(this.ecSignatureS.slice(1));
        } else {
            writer.writeBytes(this.ecSignatureS);
        }
        return writer.toArray();
    }

    // SEQUENCE {
    //   INTEGER: 00cdf3f5e083974961a9737daa2352cf08f7e652f103d4e3cc494b5ffadccf1a6d
    //   INTEGER: 69658b75c9c7523c15e6de16907350f0d0fd51114237c3e32bdd9fe92465e768
    // }
    toDer(): Uint8Array {
        let writer = new BinaryWriter();
        writer.writeNumber(0x30);
        writer.writeNumber(0x45);

        writer.writeNumber(2);
        const rFirstByte = this.ecSignatureR[0];
        writer.writeNumber(((rFirstByte >= 0x80) ? 1 : 0) + this.ecSignatureR.byteLength);
        if (rFirstByte >= 0x80) {
            writer.writeNumber(0);
        }
        writer.writeBytes(this.ecSignatureR);

        writer.writeNumber(2);
        const sFirstByte = this.ecSignatureS[0];
        writer.writeNumber(((sFirstByte >= 0x80) ? 1 : 0) + this.ecSignatureS.byteLength);
        if (sFirstByte >= 0x80) {
            writer.writeNumber(0);
        }
        writer.writeBytes(this.ecSignatureS);
        return writer.toArray();
    }
}

class SshPublicKey {
    signatureAlgorithm: string;
    algorithm: string;
    publicKeyPoint: Uint8Array;

    constructor(signatureAlgorithm: string, algorithm: string, publicKeyPoint: Uint8Array) {
        this.signatureAlgorithm = signatureAlgorithm;
        this.algorithm = algorithm;
        this.publicKeyPoint = publicKeyPoint;
    }

    static parse(buffer: Uint8Array): SshPublicKey {
        let reader = new BinaryReader(buffer);
        const signatureAlgorithm = reader.readStringToString();
        const curveAlgorithm = reader.readStringToString();
        const publicKeyPoint = reader.readString();
        if (signatureAlgorithm !== `ecdsa-sha2-${curveAlgorithm}`) {
            throw `Not supported signature algorithm ${signatureAlgorithm} or curve algorithm ${curveAlgorithm}`;
        }
        return new SshPublicKey(signatureAlgorithm, curveAlgorithm, publicKeyPoint);
    }

    toDer(): Uint8Array {
        let writer = new BinaryWriter();
        if (this.algorithm === "nistp256") {
            writer.writeBytes(decodeHex("3059301306072a8648ce3d020106082a8648ce3d030107034200"));
        } else {
            writer.writeBytes(decodeHex("3076301006072a8648ce3d020106052b81040022036200"));
        }
        writer.writeBytes(this.asPoint());
        return writer.toArray();
    }

    asPoint(): Uint8Array {
        return this.publicKeyPoint;
    }

    async importJwk(): Promise<CryptoKey> {
        let jwk = this.toJwk();
        return await crypto.subtle.importKey(
            "jwk",
            jwk,
            {
                name: "ECDSA",
                namedCurve: (this.algorithm === "nistp256") ? "P-256" : "P-384",
            },
            true,
            ["verify"],
        );
    }

    toJwk(): any {
        if (this.publicKeyPoint[0] !== 0x04) {
            throw `Invalid EC public key point: ${encodeHex(this.publicKeyPoint)}`;
        }
        let coordinateLength;
        if (this.algorithm === "nistp256") {
            coordinateLength = 256 / 8;
        } else if (this.algorithm === "nistp384") {
            coordinateLength = 384 / 8;
        } else {
            throw `Not supported algorithm: ${this.algorithm}`;
        }
        const x = this.publicKeyPoint.slice(1, coordinateLength + 1);
        const y = this.publicKeyPoint.slice(coordinateLength + 1, coordinateLength + coordinateLength + 1);
        return {
            crv: (this.algorithm === "nistp256") ? "P-256" : "P-384",
            ext: true,
            key_ops: ["verify"],
            kty: "EC",
            x: encodeBase64Url(x),
            y: encodeBase64Url(y),
        };
    }
}

function parsePemToArray(pem: string): Uint8Array {
    const lines = pem.split("\n");
    let isInPem = false;
    let innerPem = [];
    for (let i = 0; i < lines.length; i++) {
        const line = lines[i];
        if (isInPem) {
            if (line.indexOf("-----END ") === 0) {
                isInPem = false;
                break;
            } else {
                innerPem.push(line);
            }
        } else if (line.indexOf("-----BEGIN ") === 0) {
            isInPem = true;
        }
    }
    if (innerPem.length === 0) {
        return null;
    }
    return decodeBase64(innerPem.join(""));
}

async function digestString(data: string, algorithm: string): Promise<Uint8Array> {
    let hashAlgorithm = normalizeHashAlgorithm(algorithm);
    const messageBuffer = new TextEncoder().encode(data);
    const hashBuffer = await crypto.subtle.digest(hashAlgorithm, messageBuffer);
    return new Uint8Array(hashBuffer);
}

async function digestFile(filename: string, algorithm: string): Promise<Uint8Array> {
    let hashAlgorithm = normalizeHashAlgorithm(algorithm);
    const file = await Deno.open(filename, {read: true});
    const readableStream = file.readable;
    const hashBuffer = await crypto.subtle.digest(hashAlgorithm, readableStream);
    return new Uint8Array(hashBuffer);
}

function normalizeHashAlgorithm(algorithm: string): string {
    let hashAlgorithm;
    switch (algorithm.toLowerCase()) {
        case "sha256":
        case "sha-256":
            hashAlgorithm = "SHA-256";
            break;
        case "sha384":
        case "sha-384":
            hashAlgorithm = "SHA-384";
            break;
        case "sha512":
        case "sha-512":
            hashAlgorithm = "SHA-512";
            break;
        default:
            throw `Unknown algorithm: ${algorithm}`;
    }
    return hashAlgorithm;
}

async function testSshSig() {
    const TEST_SIG = `-----BEGIN SSH SIGNATURE-----
U1NIU0lHAAAAAQAAAGgAAAATZWNkc2Etc2hhMi1uaXN0cDI1NgAAAAhuaXN0cDI1
NgAAAEEE3T7r2QbJzwCwjsKfftYYBNHMHRNS2SV7YoGR4I/DcXxPrjKYzVxIKc7I
vzqUbn22C3hX4Sh/aguuaz8jQvAH0AAAAARmaWxlAAAAAAAAAAZzaGE1MTIAAABk
AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAABJAAAAIQDsOkCrkDFdGUeKYJ6smchz
qE04Ba+c5O/vN5Lld3pjFgAAACBZt+dnGiirCy/3XjKcfQYY8drtmam1QaTelyHV
WRxFFw==
-----END SSH SIGNATURE-----`;

    const sshSignature = SshSignature.parsePem(TEST_SIG);
    const data = new TextDecoder(ENCODING_UTF8).decode(decodeBase64("aGVsbG8gaGF0dGVyIDIwMjUK"));
    console.log(await sshSignature.verifyString(data));
}

// await testSshSig();