package main

import (
	"crypto/rsa"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"time"
	"os"
	"math/big"
	"fmt"
	"net"

	"github.com/alexflint/go-arg"
)

func run(args *Arguments) error {
	switch args.Type {
	case "root":
		return selfSign(args.Name)
	case "server":
		if args.Signing == "" {
			return fmt.Errorf("Signing key required")
		}
		return serverCert(args.Name, args.Signing)
	case "client":
		if args.Signing == "" {
			return fmt.Errorf("Signing key required")
		}
		return clientCert(args.Name, args.Signing)

	case "verify":
		//if args.Signing == "" {
		//	return fmt.Errorf("root cert required")
		//}
		return verify(args.Name, args.Signing)
	}

	return fmt.Errorf("unknown type: %s", args.Type)
}

func verify(target, root string) error {
	rawcert, err := os.ReadFile(target+".pem")
	if err != nil {
		return err
	}

	block, _ := pem.Decode(rawcert)
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return err
	}

	fmt.Printf("cert pubkey type: %T\n", cert.PublicKey)
	switch k := cert.PublicKey.(type) {
	case *rsa.PublicKey:
		size := k.N.BitLen();
		pub := cert.PublicKey.(*rsa.PublicKey)
		fmt.Printf("pubkey size: %d (%d)\n", pub.Size(), size)
	default:
	}

	pool := x509.NewCertPool()
	//rawroot, err := os.ReadFile(root+".pem")
	//if err != nil {
	//	return err
	//}

	//block, _ := pem.Decode(rawroot)
	//rootcert, err := x509.ParseCertificate(block.Bytes)
	//if err != nil {
	//	return err
	//}
	if root != "" {
		fmt.Println("loading "+root+".pem")
		rootRaw, err := os.ReadFile(root+".pem")
		if err != nil {
			return err
		}

		if !pool.AppendCertsFromPEM(rootRaw) {
			return fmt.Errorf("add pem not ok")
		}
	}

	fmt.Printf("primary: %d (%08b)\n", cert.KeyUsage, cert.KeyUsage)
	for _, usage := range cert.ExtKeyUsage {
		fmt.Printf("> %d (%08b)\n", usage, usage)
	}

	chains, err := cert.Verify(x509.VerifyOptions{
		Roots: pool,
		KeyUsages: []x509.ExtKeyUsage{
			x509.ExtKeyUsageClientAuth,
		},
	})
	fmt.Println("len(chains):", len(chains))
	if err != nil {
		return err
	}

	fmt.Println("verify OK")
	return nil
}

func selfSign(name string) error {
	notBefore := time.Now()
	notAfter := notBefore.Add(90*24*time.Hour)

	/*
		Root Cert
	*/
	rootkey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return err
	}

	err = os.WriteFile(name+".key", pem.EncodeToMemory(&pem.Block{Type:"RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(rootkey)}), 0700)
	if err != nil {
		return err
	}

	rootTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization: []string{"Zorchenhimer LLC"},
			CommonName: "Root CA",
		},
		NotBefore: notBefore,
		NotAfter: notAfter,
		KeyUsage: x509.KeyUsageCertSign,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA: true,
	}

	rootcert, err := x509.CreateCertificate(rand.Reader, rootTemplate, rootTemplate, &rootkey.PublicKey, rootkey)
	if err != nil {
		return err
	}

	err = os.WriteFile(name+".pem", pem.EncodeToMemory(
		&pem.Block{
			Type:"CERTIFICATE",
			Bytes: rootcert,
		}), 0700)
	if err != nil {
		return err
	}

	return nil
}

func loadRoot(name string) (*x509.Certificate, *pem.Block, *rsa.PrivateKey, error) {
	rawroot, err := os.ReadFile(name+".pem")
	if err != nil {
		return nil, nil, nil, err
	}

	rootpem, _ := pem.Decode(rawroot)
	rootcert, err := x509.ParseCertificate(rootpem.Bytes)
	if err != nil {
		return nil, nil, nil, err
	}

	rawkey, err := os.ReadFile(name+".key")
	if err != nil {
		return nil, nil, nil, err
	}

	block, _ := pem.Decode(rawkey)
	rootkey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, nil, nil, err
	}

	return rootcert, rootpem, rootkey, nil
}

func serverCert(name, rootname string) error {
	rootcert, _, rootkey, err := loadRoot(rootname)
	if err != nil {
		return err
	}
	notBefore := time.Now()
	notAfter := notBefore.Add(90*24*time.Hour)

	/*
		Server Cert
	*/

	serverkey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return err
	}

	err = os.WriteFile(name+".key",
		pem.EncodeToMemory(&pem.Block{
			Type:"RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(serverkey),
		}), 0700)
	if err != nil {
		return err
	}

	serverTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject: pkix.Name{
			Organization: []string{"Zorchenhimer LLC"},
			CommonName: "Server Cert",
		},
		NotBefore: notBefore,
		NotAfter: notAfter,
		KeyUsage: x509.KeyUsageDigitalSignature |
				  x509.KeyUsageKeyEncipherment,
		ExtKeyUsage: []x509.ExtKeyUsage{
			x509.ExtKeyUsageServerAuth,
			x509.ExtKeyUsageClientAuth,
		},
		BasicConstraintsValid: true,
		IsCA: false,

		DNSNames: []string{
			"localhost",
			"ryuko.local",
		},
		IPAddresses: []net.IP{
			net.ParseIP("127.0.0.1"),
			net.ParseIP("192.168.1.10"),
		},
	}

	servercert, err := x509.CreateCertificate(
		rand.Reader, serverTemplate, rootcert, &serverkey.PublicKey, rootkey)
	if err != nil {
		return err
	}

	err = os.WriteFile(name+".pem",
		pem.EncodeToMemory(&pem.Block{Type:"CERTIFICATE", Bytes: servercert}), 0700)
	if err != nil {
		return err
	}

	return nil
}

func clientCert(name, root string) error {
	rootcert, rootblock, rootkey, err := loadRoot(root)
	if err != nil {
		return err
	}
	notBefore := time.Now()
	notAfter := notBefore.Add(90*24*time.Hour)

	/*
		Client Cert
	*/

	clientkey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return err
	}

	err = os.WriteFile(name+".key", pem.EncodeToMemory(
		&pem.Block{
			Type:"RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(clientkey),
		}), 0700)
	if err != nil {
		return err
	}

	clientTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(3),
		Subject: pkix.Name{
			Organization: []string{"Zorchenhimer LLC"},
			CommonName: "Client Cert",
		},
		NotBefore: notBefore,
		NotAfter: notAfter,
		//KeyUsage: x509.KeyUsageDigitalSignature |
		//          x509.KeyUsageDataEncipherment |
		//          x509.KeyUsageKeyEncipherment |
		//		  x509.KeyUsageContentCommitment |
		//		  x509.KeyUsageKeyAgreement,
		KeyUsage: x509.KeyUsageKeyEncipherment |
				  x509.KeyUsageDigitalSignature,
		ExtKeyUsage: []x509.ExtKeyUsage{
			//x509.ExtKeyUsageAny,
			x509.ExtKeyUsageClientAuth,
		},
		BasicConstraintsValid: true,
		IsCA: false,
		IPAddresses: []net.IP{
			net.ParseIP("127.0.0.1"),
		},
	}

	clientcert, err := x509.CreateCertificate(
		rand.Reader, clientTemplate, rootcert, &clientkey.PublicKey, rootkey)
	if err != nil {
		return err
	}

	file, err := os.Create(name+".pem")
	if err != nil {
		return err
	}
	defer file.Close()

	err = pem.Encode(file, &pem.Block{
		Type:"CERTIFICATE",
		Bytes: clientcert,
	})
	if err != nil {
		return err
	}

	//err = pem.Encode(file, &pem.Block{
	//	Type: "PUBLIC KEY",
	//	Bytes: clientkey.Pubkey
	//})

	err = pem.Encode(file, rootblock)
	if err != nil {
		return err
	}

	//err = os.WriteFile(name+".pem",
	//	pem.EncodeToMemory(&pem.Block{Type:"CERTIFICATE", Bytes: clientcert}), 0700)
	//if err != nil {
	//	return err
	//}

	return nil
}

type Arguments struct {
	Type string `arg:"positional,required"`
	Name string `arg:"positional,required"`
	Signing string `arg:"positional"`
}

func main() {
	args := &Arguments{}
	arg.MustParse(args)

	err := run(args)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}