/*
 * Copyright 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 *
 */

package harmonyos

import (
	"crypto/ecdsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"github.com/cristalhq/jwt/v3"
	"io/ioutil"
	"net/http"
	"strings"
)

const (
	CA_CERT_FILE_PATH = "manifest/config/RootCaG2Ecdsa.cer" // TODO: Need to replace it with the actual value.

	LEAF_CERT_OID = "1.3.6.1.4.1.2011.2.415.1.1"

	// HEADER_PARAM_ALG Algorithm type. The value is always ES256.
	HEADER_PARAM_ALG = "ES256"

	X5C_CHAIN_LENGTH = 3
)

type JWSChecker struct {
}

type Header struct {
	// Alg Algorithm type. The value is always ES256.
	Alg string `json:"alg"`
	// X5c Indicates the X.509 certificate chain. The sequence is leaf certificate, intermediate certificate, and root certificate.
	X5c []string `json:"x5c"`
}

func (jwsChecker *JWSChecker) CheckAndDecodeJWS(jwsStr string) (string, error) {
	jwsStrArr := strings.Split(jwsStr, ".")
	if len(jwsStrArr) != X5C_CHAIN_LENGTH {
		// TODO: Need to replace it with the actual business logic.
		return "", errors.New("invalid jwsStr length")
	}

	certChain, err := jwsChecker.extractCertChain(jwsStrArr)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return "", err
	}

	rootCert, err := jwsChecker.loadRootCert()
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return "", err
	}

	if err = jwsChecker.verifyCertChain(certChain, rootCert); err != nil {
		// TODO: Need to replace it with the actual business logic.
		return "", err
	}

	if err = jwsChecker.verifyOID(certChain[0]); err != nil {
		// TODO: Need to replace it with the actual business logic.
		return "", err
	}

	if err = jwsChecker.verifyCRLFromCerts(certChain[:2]); err != nil {
		// TODO: Determine whether to ignore CRL certificate download exception based on application Security Policy.
		return "", err
	}

	if err = jwsChecker.verifySignature(jwsStr, certChain[0], string(jwt.ES256)); err != nil {
		// TODO: Need to replace it with the actual business logic.
		return "", err
	}

	payLoadBytes, err := base64.RawStdEncoding.DecodeString(jwsStrArr[1])
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return "", err
	}

	return string(payLoadBytes), err
}

func (jwsChecker *JWSChecker) loadRootCert() (*x509.Certificate, error) {
	rootCABytes, err := ioutil.ReadFile(CA_CERT_FILE_PATH)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return nil, err
	}

	rootCerBlock, _ := pem.Decode(rootCABytes)
	if rootCerBlock == nil {
		// TODO: Need to replace it with the actual business logic.
		return nil, err
	}

	rootCer, err := x509.ParseCertificate(rootCerBlock.Bytes)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return nil, err
	}

	return rootCer, err
}

func (jwsChecker *JWSChecker) extractCertChain(jwsStrArr []string) ([]*x509.Certificate, error) {
	headerByte, err := base64.RawStdEncoding.DecodeString(jwsStrArr[0])
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return nil, err
	}

	var header Header
	err = json.Unmarshal(headerByte, &header)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return nil, err
	}

	certChain := make([]*x509.Certificate, len(header.X5c))

	for i := range header.X5c {
		certByte, err := base64.StdEncoding.DecodeString(header.X5c[i])
		if err != nil {
			// TODO: Need to replace it with the actual business logic.
			return nil, err
		}
		cert, err := x509.ParseCertificate(certByte)
		if err != nil {
			// TODO: Need to replace it with the actual business logic.
			return nil, err
		}
		certChain[i] = cert
	}

	return certChain, err
}

func (jwsChecker *JWSChecker) verifyCertChain(certChain []*x509.Certificate, rootCert *x509.Certificate) error {
	roots := x509.NewCertPool()
	roots.AddCert(rootCert)

	intermedia := x509.NewCertPool()
	intermedia.AddCert(certChain[1])

	opts := x509.VerifyOptions{
		Roots:         roots,
		Intermediates: intermedia,
	}
	_, err := rootCert.Verify(opts)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return err
	}

	_, err = certChain[0].Verify(opts)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return err
	}

	return err
}

func (jwsChecker *JWSChecker) verifyCRLFromCerts(certChain []*x509.Certificate) error {
	for _, cert := range certChain {
		err := checkCrl(cert)
		if err != nil {
			// TODO: Determine whether to ignore CRL certificate download exception based on application Security Policy.
			return err
		}
	}
	return nil
}

func checkCrl(cert *x509.Certificate) error {
	if len(cert.CRLDistributionPoints) == 0 {
		return errors.New("No CRL distribution points found in certificate")
	}

	resp, err := http.Get(cert.CRLDistributionPoints[0])
	if err != nil {
		// TODO: Determine whether to ignore CRL certificate download exception based on application Security Policy.
		return err
	}
	defer resp.Body.Close()

	crlData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return err
	}

	crl, err := x509.ParseCRL(crlData)
	if err != nil {
		// TODO: Need to replace it with the actual business logic.
		return err
	}

	if crl.TBSCertList.RevokedCertificates == nil {
		return nil
	}

	for _, revokedCert := range crl.TBSCertList.RevokedCertificates {
		if revokedCert.SerialNumber.Cmp(cert.SerialNumber) == 0 {
			// TODO: Need to replace it with the actual business logic.
			return fmt.Errorf("certificate %s has been revoked", cert.Subject.CommonName)
		}
	}

	return nil
}

func (jwsChecker *JWSChecker) verifyOID(leafCert *x509.Certificate) error {
	var isExist = false
	for i := range leafCert.Extensions {
		if LEAF_CERT_OID == leafCert.Extensions[i].Id.String() && leafCert.Extensions[i].Critical == false {
			isExist = true
		}
	}

	if !isExist {
		// TODO: Need to replace it with the actual business logic.
		return errors.New("oid verification failed")
	}

	return nil
}

func (jwsChecker *JWSChecker) verifySignature(jwsStr string, leafCert *x509.Certificate, algParam string) error {
	if algParam != jwt.ES256.String() {
		// TODO: Need to replace it with the actual business logic.
		return errors.New("invalid algorithm")
	}
	switch pubicKey := leafCert.PublicKey.(type) {
	case *ecdsa.PublicKey:
		verifier, err := jwt.NewVerifierES(jwt.ES256, pubicKey)
		if err != nil {
			return err
		}

		token := []byte(jwsStr)
		_, err = jwt.ParseAndVerify(token, verifier)
		if err != nil {
			return err
		}
	default:
		// TODO: Need to replace it with the actual business logic.
		return errors.New("public key not supported")
	}
	return nil
}