package main

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"github.com/urfave/cli/v2"
)

type ErrorResponse struct {
	Success bool   `json:"success"`
	Error   string `json:"error"`
}

type ExternalSpecResponse struct {
	Success       bool     `json:"success"`
	Agent         string   `json:"agent"`
	Specification string   `json:"specification"`
	Commands      []string `json:"commands"`
}

type ExternalPublicKeyResponse struct {
	Success         bool   `json:"success"`
	PublicKeyBase64 string `json:"public_key_base64"`
}

type ExternalSignatureResponse struct {
	Success         bool   `json:"success"`
	SignatureBase64 string `json:"signature_base64"`
}

func printResponse(response any) error {
	errorMessageJsonBytes, jsonErr := json.Marshal(response)
	if jsonErr != nil {
		return jsonErr
	} else {
		fmt.Println(string(errorMessageJsonBytes))
		return nil
	}
}

func printResponseSlient(response any) {
	if err := printResponse(response); err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "%v\n", err)
	}
}

func main() {
	err := innerMain()
	if err != nil {
		errorMessage := &ErrorResponse{
			Success: false,
			Error:   fmt.Sprintf("%v", err),
		}
		if printResponse(errorMessage) != nil {
			_, _ = fmt.Fprintf(os.Stderr, "%v\n", err)
		}
		os.Exit(1)
	}
}

func innerMain() error {
	app := cli.App{
		Name:  "external-signer-soft",
		Usage: "External signer software",
		Commands: []*cli.Command{
			buildGenerateKeypair(),
			buildExternalSpecCommand(),
			buildExternalPublicKeyCommand(),
			buildExternalSignCommand(),
		},
		Action: func(ctx *cli.Context) error {
			fmt.Println("External signer software, specification: https://openwebstandard.org/rfc1")
			return nil
		},
	}

	if err := app.Run(os.Args); err != nil {
		return err
	}
	return nil
}

func buildGenerateKeypair() *cli.Command {
	return &cli.Command{
		Name:  "generate_keypair",
		Usage: "Generate keypair",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "type",
				Usage:    "Type, enums: rsa2048, rsa3072, rsa4096, p256, p384, p521",
				Required: true,
			},
		},
		Action: func(ctx *cli.Context) error {
			keyType := ctx.String("type")
			privateKeyDer, err := generateKeypair(keyType)
			if err != nil {
				return err
			}

			fmt.Printf("Private Key:\n%s\n", base64.StdEncoding.EncodeToString(privateKeyDer))

			return nil
		},
	}
}

func generateKeypair(keyType string) ([]byte, error) {
	keyType = strings.ToLower(keyType)
	if strings.HasPrefix(keyType, "rsa") {
		var bitLen int
		if keyType == "rsa2048" {
			bitLen = 2048
		} else if keyType == "rsa3072" {
			bitLen = 3072
		} else if keyType == "rsa4096" {
			bitLen = 4096
		} else {
			return nil, fmt.Errorf("bad key type: %s", keyType)
		}
		privateKey, err := rsa.GenerateKey(rand.Reader, bitLen)
		if err != nil {
			return nil, err
		}
		return x509.MarshalPKCS8PrivateKey(privateKey)
	}
	var curve elliptic.Curve
	if keyType == "p256" {
		curve = elliptic.P256()
	} else if keyType == "p384" {
		curve = elliptic.P384()
	} else if keyType == "p521" {
		curve = elliptic.P521()
	} else {
		return nil, fmt.Errorf("bad key type: %s", keyType)
	}
	privateKey, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err != nil {
		return nil, err
	}
	return x509.MarshalPKCS8PrivateKey(privateKey)
}

func buildExternalSpecCommand() *cli.Command {
	return &cli.Command{
		Name:  "external_spec",
		Usage: "External specification",
		Flags: []cli.Flag{},
		Action: func(ctx *cli.Context) error {
			externalSpecResponse := &ExternalSpecResponse{
				Success:       true,
				Agent:         "external-signer-soft/0.1.0",
				Specification: "External/1.0.0-alpha",
				Commands: []string{
					"external_public_key",
					"external_sign",
				},
			}
			printResponseSlient(externalSpecResponse)
			return nil
		},
	}
}

func buildExternalPublicKeyCommand() *cli.Command {
	return &cli.Command{
		Name:  "external_public_key",
		Usage: "External public key",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "parameter",
				Usage:    "Parameter",
				Required: true,
			},
		},
		Action: func(ctx *cli.Context) error {
			parameter := ctx.String("parameter")
			privateKeyDer, err := base64.StdEncoding.DecodeString(parameter)
			if err != nil {
				return err
			}
			privateKey, err := x509.ParsePKCS8PrivateKey(privateKeyDer)
			if err != nil {
				return err
			}

			var publicKeyDer []byte
			var publicKeyErr error
			if rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey); ok {
				publicKey := rsaPrivateKey.PublicKey
				publicKeyDer, publicKeyErr = x509.MarshalPKIXPublicKey(&publicKey)
			} else if ecdsaPrivateKey, ok := privateKey.(*ecdsa.PrivateKey); ok {
				publicKey := ecdsaPrivateKey.PublicKey
				publicKeyDer, publicKeyErr = x509.MarshalPKIXPublicKey(&publicKey)
			} else {
				return fmt.Errorf("bad paramater")
			}

			if publicKeyErr != nil {
				return publicKeyErr
			}

			externalPublicKeyResponse := &ExternalPublicKeyResponse{
				Success:         true,
				PublicKeyBase64: base64.StdEncoding.EncodeToString(publicKeyDer),
			}

			printResponseSlient(externalPublicKeyResponse)
			return nil
		},
	}
}

func buildExternalSignCommand() *cli.Command {
	return &cli.Command{
		Name:  "external_sign",
		Usage: "External sign",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:     "parameter",
				Usage:    "Parameter",
				Required: true,
			},
			&cli.StringFlag{
				Name:     "alg",
				Usage:    "Algorithm",
				Required: true,
			},
			&cli.StringFlag{
				Name:     "message-base64",
				Usage:    "Message base64 encoded",
				Required: true,
			},
			&cli.StringFlag{
				Name:  "message-type",
				Usage: "Message type raw(default), sha256, sha384 or sha512",
			},
		},
		Action: func(ctx *cli.Context) error {
			parameter := ctx.String("parameter")
			alg := ctx.String("alg")
			messageBase64 := ctx.String("message-base64")
			messageType := ctx.String("message-type")
			message, err := base64.StdEncoding.DecodeString(messageBase64)
			if err != nil {
				return err
			}

			privateKeyDer, err := base64.StdEncoding.DecodeString(parameter)
			if err != nil {
				return err
			}
			privateKey, err := x509.ParsePKCS8PrivateKey(privateKeyDer)
			if err != nil {
				return err
			}

			var digest []byte
			var hash crypto.Hash
			expectedMessageType := ""
			if alg == "ES256" || alg == "RS256" {
				expectedMessageType = "sha256"
				digest32 := sha256.Sum256(message)
				digest = digest32[:]
				hash = crypto.SHA256
			} else if alg == "ES384" || alg == "RS384" {
				expectedMessageType = "sha384"
				sha384 := crypto.SHA384.New()
				digest = sha384.Sum(message)
				hash = crypto.SHA384
			} else if alg == "ES512" || alg == "RS512" {
				expectedMessageType = "sha512"
				digest64 := sha512.Sum512(message)
				digest = digest64[:]
				hash = crypto.SHA512
			} else {
				return fmt.Errorf("invalid algorithm: %s", alg)
			}
			if messageType != "" && messageType != "raw" {
				if messageType != expectedMessageType {
					return fmt.Errorf("invalid message type %s vs %s", messageType, alg)
				}
				digest = message[:]
			}

			var signature []byte
			var signatureErr error
			if rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey); ok {
				if !strings.HasPrefix(alg, "RS") {
					return fmt.Errorf("invalid algorithm for RSA: %s", alg)
				}
				signature, signatureErr = rsaPrivateKey.Sign(rand.Reader, digest, hash)
			} else if ecdsaPrivateKey, ok := privateKey.(*ecdsa.PrivateKey); ok {
				paramsName := ecdsaPrivateKey.Params().Name
				if paramsName == "P-256" && alg != "ES256" {
					return fmt.Errorf("invalid algorithm for P256: %s", alg)
				}
				if paramsName == "P-384" && alg != "ES384" {
					return fmt.Errorf("invalid algorithm for P384: %s", alg)
				}
				if paramsName == "P-521" && alg != "ES521" {
					return fmt.Errorf("invalid algorithm for P521: %s", alg)
				}
				signature, signatureErr = ecdsaPrivateKey.Sign(rand.Reader, digest, hash)
			} else {
				return fmt.Errorf("bad paramater")
			}
			if signatureErr != nil {
				return fmt.Errorf("sign with private key failed: %v", signatureErr)
			}

			externalSignatureResponse := ExternalSignatureResponse{
				Success:         true,
				SignatureBase64: base64.StdEncoding.EncodeToString(signature),
			}

			printResponseSlient(externalSignatureResponse)
			return nil
		},
	}
}