// GenerateKeyPair generates a private/public key pair,
// keys are returned as hex-encoded strings
func GenerateKeyPair() (private_key_hex, public_key_hex string) {
	// generate keys
	private_key, err := ecdsa.GenerateKey(elliptic.P224(), rand.Reader)
	if err != nil {
		panic(err)
	}

	// marshal private key
	private_key_bytes, err := x509.MarshalECPrivateKey(private_key)
	if err != nil {
		panic(err)
	}

	// marshal public key
	public_key_bytes, err := x509.MarshalPKIXPublicKey(&private_key.PublicKey)
	if err != nil {
		panic(err)
	}

	// hex encode and return result
	private_key_hex = hex.EncodeToString(private_key_bytes)
	public_key_hex = hex.EncodeToString(public_key_bytes)

	return private_key_hex, public_key_hex
}

// generateFromTemplate generates a certificate from the given template and signed by
// the given parent, storing the results in a certificate and key file.
func generateFromTemplate(certFile, keyFile string, template, parent *x509.Certificate, key crypto.PrivateKey, parentKey crypto.PrivateKey) error {
	derBytes, err := x509.CreateCertificate(rand.Reader, template, parent, key.(crypto.Signer).Public(), parentKey)
	if err != nil {
		return err
	}

	certOut, err := os.Create(certFile)
	if err != nil {
		return err
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()

	keyOut, err := os.OpenFile(keyFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	defer keyOut.Close()

	switch v := key.(type) {
	case *rsa.PrivateKey:
		keyBytes := x509.MarshalPKCS1PrivateKey(v)
		pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: keyBytes})
	case *ecdsa.PrivateKey:
		keyBytes, err := x509.MarshalECPrivateKey(v)
		if err != nil {
			return err
		}
		pem.Encode(keyOut, &pem.Block{Type: "EC PRIVATE KEY", Bytes: keyBytes})
	default:
		return fmt.Errorf("Unsupport private key type: %#v", key)
	}

	return nil
}

func (ctrl *Controller) AddHostKeyFromCluster(host string) error {
	generate := func() (string, error) {
		// Generate an ECDSA key.
		key, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
		if err != nil {
			return "", err
		}
		derBytes, err := x509.MarshalECPrivateKey(key)
		if err != nil {
			return "", err
		}
		w := &bytes.Buffer{}
		if err := pem.Encode(w, &pem.Block{Type: "EC PRIVATE KEY", Bytes: derBytes}); err != nil {
			return "", err
		}
		return w.String(), nil
	}
	pemString, err := ctrl.cluster.GetValueWithDefault(fmt.Sprintf("console/%s", host), generate)
	if err != nil {
		return fmt.Errorf("failed to get/generate host key: %s", err)
	}

	signer, err := ssh.ParsePrivateKey([]byte(pemString))
	if err != nil {
		return fmt.Errorf("failed to parse host key: %s", err)
	}

	ctrl.config.AddHostKey(signer)
	return nil
}

// PrivateKeyToEncryptedPEM converts a private key to an encrypted PEM
func PrivateKeyToEncryptedPEM(privateKey interface{}, pwd []byte) ([]byte, error) {
	switch k := privateKey.(type) {
	case *ecdsa.PrivateKey:
		if k == nil {
			return nil, errors.New("Invalid ecdsa private key. It must be different from nil.")
		}

		raw, err := x509.MarshalECPrivateKey(k)

		if err != nil {
			return nil, err
		}

		block, err := x509.EncryptPEMBlock(
			rand.Reader,
			"ECDSA PRIVATE KEY",
			raw,
			pwd,
			x509.PEMCipherAES256)

		if err != nil {
			return nil, err
		}

		return pem.EncodeToMemory(block), nil

	default:
		return nil, errors.New("Invalid key type. It must be *ecdsa.PrivateKey")
	}
}

// StartKEX prepares a new key exchange. It returns an initialised
// session handle and a signed public key that should be sent to the
// peer. peer and FinishKEX called to finalise the session. It returns a
// new session handle and a signed public key that should be sent to the
// peer.  The returned session handle has ephemeral private key data in
// it, but the shared key is not yet set up. After this call, the session
// cannot encrypt or decrypt.
func StartKEX(signer *ecdsa.PrivateKey) (*Session, []byte, error) {
	priv, err := ecdsa.GenerateKey(signer.Curve, rand.Reader)
	if err != nil {
		return nil, nil, err
	}

	skey := signedKey{}
	skey.Public, err = x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		return nil, nil, err
	}

	hashedPub := sha256.Sum256(skey.Public)

	skey.R, skey.S, err = ecdsa.Sign(rand.Reader, signer, hashedPub[:])
	if err != nil {
		return nil, nil, err
	}

	kex := &Session{}
	kex.priv, err = x509.MarshalECPrivateKey(priv)
	if err != nil {
		return nil, nil, err
	}

	out, err := asn1.Marshal(skey)
	if err != nil {
		return nil, nil, err
	}

	return kex, out, nil
}

// PrivateKeyToDER marshals a private key to der
func PrivateKeyToDER(privateKey *ecdsa.PrivateKey) ([]byte, error) {
	if privateKey == nil {
		return nil, errors.New("Invalid ecdsa private key. It must be different from nil.")
	}

	return x509.MarshalECPrivateKey(privateKey)
}

// PrivateKeyToDER marshals a private key to der
func PrivateKeyToDER(privateKey *ecdsa.PrivateKey) ([]byte, error) {
	if privateKey == nil {
		return nil, utils.ErrNilArgument
	}

	return x509.MarshalECPrivateKey(privateKey)
}

// savePrivateKey saves a PEM-encoded ECC/RSA private key to file.
func savePrivateKey(key crypto.PrivateKey, file string) error {
	var pemType string
	var keyBytes []byte
	switch key := key.(type) {
	case *ecdsa.PrivateKey:
		var err error
		pemType = "EC"
		keyBytes, err = x509.MarshalECPrivateKey(key)
		if err != nil {
			return err
		}
	case *rsa.PrivateKey:
		pemType = "RSA"
		keyBytes = x509.MarshalPKCS1PrivateKey(key)
	}

	pemKey := pem.Block{Type: pemType + " PRIVATE KEY", Bytes: keyBytes}
	keyOut, err := os.Create(file)
	if err != nil {
		return err
	}
	keyOut.Chmod(0600)
	defer keyOut.Close()
	return pem.Encode(keyOut, &pemKey)
}

// Write a private key in PEM form.
func SavePrivateKey(w io.Writer, pk crypto.PrivateKey) error {
	var kb []byte
	var hdr string
	var err error

	switch v := pk.(type) {
	case *rsa.PrivateKey:
		kb = x509.MarshalPKCS1PrivateKey(v)
		hdr = "RSA PRIVATE KEY"
	case *ecdsa.PrivateKey:
		kb, err = x509.MarshalECPrivateKey(v)
		hdr = "EC PRIVATE KEY"
	default:
		return fmt.Errorf("unsupported private key type: %T", pk)
	}
	if err != nil {
		return err
	}

	err = pem.Encode(w, &pem.Block{
		Type:  hdr,
		Bytes: kb,
	})
	if err != nil {
		return err
	}

	return nil
}

func generateCertificates(privKeyFile string, certKeyFile string) error {
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return err
	}

	sn, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	template := x509.Certificate{
		KeyUsage:     x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		SerialNumber: sn,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	rawBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
	if err != nil {
		return err
	}
	err = writeFile(certKeyFile, "CERTIFICATE", rawBytes)
	if err != nil {
		return err
	}
	privBytes, err := x509.MarshalECPrivateKey(privateKey)
	if err != nil {
		return err
	}
	err = writeFile(privKeyFile, "EC PRIVATE KEY", privBytes)
	return err
}

func CreatePrivateKey(d *schema.ResourceData, meta interface{}) error {
	keyAlgoName := d.Get("algorithm").(string)
	var keyFunc keyAlgo
	var ok bool
	if keyFunc, ok = keyAlgos[keyAlgoName]; !ok {
		return fmt.Errorf("invalid key_algorithm %#v", keyAlgoName)
	}

	key, err := keyFunc(d)
	if err != nil {
		return err
	}

	var keyPemBlock *pem.Block
	switch k := key.(type) {
	case *rsa.PrivateKey:
		keyPemBlock = &pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: x509.MarshalPKCS1PrivateKey(k),
		}
	case *ecdsa.PrivateKey:
		keyBytes, err := x509.MarshalECPrivateKey(k)
		if err != nil {
			return fmt.Errorf("error encoding key to PEM: %s", err)
		}
		keyPemBlock = &pem.Block{
			Type:  "EC PRIVATE KEY",
			Bytes: keyBytes,
		}
	default:
		return fmt.Errorf("unsupported private key type")
	}
	keyPem := string(pem.EncodeToMemory(keyPemBlock))

	pubKey := publicKey(key)
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return fmt.Errorf("failed to marshal public key: %s", err)
	}
	pubKeyPemBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	}

	d.SetId(hashForState(string((pubKeyBytes))))
	d.Set("private_key_pem", keyPem)
	d.Set("public_key_pem", string(pem.EncodeToMemory(pubKeyPemBlock)))

	sshPubKey, err := ssh.NewPublicKey(pubKey)
	if err == nil {
		// Not all EC types can be SSH keys, so we'll produce this only
		// if an appropriate type was selected.
		sshPubKeyBytes := ssh.MarshalAuthorizedKey(sshPubKey)
		d.Set("public_key_openssh", string(sshPubKeyBytes))
	} else {
		d.Set("public_key_openssh", "")
	}

	return nil
}

func marshalPKCS8PrivateKey(key interface{}) (der []byte, err error) {
	var privKey pkcs8
	switch key := key.(type) {
	case *rsa.PrivateKey:
		privKey.Algo.Algorithm = oidPublicKeyRSA
		// This is a NULL parameters value which is technically
		// superfluous, but most other code includes it.
		privKey.Algo.Parameters = asn1.RawValue{
			Tag: 5,
		}
		privKey.PrivateKey = x509.MarshalPKCS1PrivateKey(key)
	case *ecdsa.PrivateKey:
		privKey.Algo.Algorithm = oidPublicKeyECDSA
		namedCurveOID, ok := oidFromNamedCurve(key.Curve)
		if !ok {
			return nil, errors.New("pkcs12: unknown elliptic curve")
		}
		if privKey.Algo.Parameters.FullBytes, err = asn1.Marshal(namedCurveOID); err != nil {
			return nil, errors.New("pkcs12: failed to embed OID of named curve in PKCS#8: " + err.Error())
		}
		if privKey.PrivateKey, err = x509.MarshalECPrivateKey(key); err != nil {
			return nil, errors.New("pkcs12: failed to embed EC private key in PKCS#8: " + err.Error())
		}
	default:
		return nil, errors.New("pkcs12: only RSA and ECDSA private keys supported")
	}
	return asn1.Marshal(privKey)
}

// PrivateKeyToEncryptedPEM converts a private key to an encrypted PEM
func PrivateKeyToEncryptedPEM(privateKey interface{}, pwd []byte) ([]byte, error) {
	switch x := privateKey.(type) {
	case *ecdsa.PrivateKey:
		raw, err := x509.MarshalECPrivateKey(x)

		if err != nil {
			return nil, err
		}

		block, err := x509.EncryptPEMBlock(
			rand.Reader,
			"ECDSA PRIVATE KEY",
			raw,
			pwd,
			x509.PEMCipherAES256)

		if err != nil {
			return nil, err
		}

		return pem.EncodeToMemory(block), nil

	default:
		return nil, ErrInvalidKey
	}
}

// marshalPKCS8PrivateKey converts a private key to PKCS#8 encoded form.
// See http://www.rsa.com/rsalabs/node.asp?id=2130 and RFC5208.
func marshalPKCS8PrivateKey(key interface{}) (der []byte, err error) {
	pkcs := pkcs8{
		Version: 0,
	}

	switch key := key.(type) {
	case *rsa.PrivateKey:
		pkcs.Algo = pkix.AlgorithmIdentifier{
			Algorithm:  oidPublicKeyRSA,
			Parameters: nullAsn,
		}
		pkcs.PrivateKey = x509.MarshalPKCS1PrivateKey(key)
	case *ecdsa.PrivateKey:
		bytes, err := x509.MarshalECPrivateKey(key)
		if err != nil {
			return nil, errors.New("x509: failed to marshal to PKCS#8: " + err.Error())
		}

		pkcs.Algo = pkix.AlgorithmIdentifier{
			Algorithm:  oidPublicKeyECDSA,
			Parameters: nullAsn,
		}
		pkcs.PrivateKey = bytes
	default:
		return nil, errors.New("x509: PKCS#8 only RSA and ECDSA private keys supported")
	}

	bytes, err := asn1.Marshal(pkcs)
	if err != nil {
		return nil, errors.New("x509: failed to marshal to PKCS#8: " + err.Error())
	}

	return bytes, nil
}

func (ca *CA) createCAKeyPair(name string) *ecdsa.PrivateKey {
	caLogger.Debug("Creating CA key pair.")

	curve := primitives.GetDefaultCurve()

	priv, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err == nil {
		raw, _ := x509.MarshalECPrivateKey(priv)
		cooked := pem.EncodeToMemory(
			&pem.Block{
				Type:  "ECDSA PRIVATE KEY",
				Bytes: raw,
			})
		err = ioutil.WriteFile(ca.path+"/"+name+".priv", cooked, 0644)
		if err != nil {
			caLogger.Panic(err)
		}

		raw, _ = x509.MarshalPKIXPublicKey(&priv.PublicKey)
		cooked = pem.EncodeToMemory(
			&pem.Block{
				Type:  "ECDSA PUBLIC KEY",
				Bytes: raw,
			})
		err = ioutil.WriteFile(ca.path+"/"+name+".pub", cooked, 0644)
		if err != nil {
			caLogger.Panic(err)
		}
	}
	if err != nil {
		caLogger.Panic(err)
	}

	return priv
}

func TestECDSAVerifierOtherCurves(t *testing.T) {
	curves := []elliptic.Curve{elliptic.P256(), elliptic.P384(), elliptic.P521()}

	for _, curve := range curves {
		ecdsaPrivKey, err := ecdsa.GenerateKey(curve, rand.Reader)

		// Get a DER-encoded representation of the PublicKey
		ecdsaPubBytes, err := x509.MarshalPKIXPublicKey(&ecdsaPrivKey.PublicKey)
		assert.NoError(t, err, "failed to marshal public key")

		// Get a DER-encoded representation of the PrivateKey
		ecdsaPrivKeyBytes, err := x509.MarshalECPrivateKey(ecdsaPrivKey)
		assert.NoError(t, err, "failed to marshal private key")

		testECDSAPubKey := data.NewECDSAPublicKey(ecdsaPubBytes)
		testECDSAKey, err := data.NewECDSAPrivateKey(testECDSAPubKey, ecdsaPrivKeyBytes)
		assert.NoError(t, err, "failed to read private key")

		// Sign some data using ECDSA
		message := []byte("test data for signing")
		hashed := sha256.Sum256(message)
		signedData, err := ecdsaSign(testECDSAKey, hashed[:])
		assert.NoError(t, err)

		// Create and call Verify on the verifier
		ecdsaVerifier := ECDSAVerifier{}
		err = ecdsaVerifier.Verify(testECDSAKey, signedData, message)
		assert.NoError(t, err, "expecting success but got error while verifying data using ECDSA")

		// Make sure an invalid signature fails verification
		signedData[0]++
		err = ecdsaVerifier.Verify(testECDSAKey, signedData, message)
		assert.Error(t, err, "expecting error but got success while verifying data using ECDSA")
	}
}

func marshalKey(key *ecdsa.PrivateKey) ([]byte, error) {
	data, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, err
	}
	return pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: data}), nil
}

// GenerateKeys returns a new key pair, with the private and public key
// encoded in PEM format.
func GenerateKeys() (privKey []byte, pubKey []byte, err error) {
	// Generate a new key pair
	key, err := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)
	if err != nil {
		return nil, nil, err
	}

	// Marshal the private key
	bs, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, nil, err
	}

	// Encode it in PEM format
	privKey = pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: bs,
	})

	// Marshal the public key
	bs, err = x509.MarshalPKIXPublicKey(key.Public())
	if err != nil {
		return nil, nil, err
	}

	// Encode it in PEM format
	pubKey = pem.EncodeToMemory(&pem.Block{
		Type:  "EC PUBLIC KEY",
		Bytes: bs,
	})

	return
}

func TestTLSConnection(t *testing.T) {
	reactor := NewReactor()
	client := reactor.CreateServer("local")

	initialiseServerConnection(client)

	// generate a test certificate to use
	priv, _ := ecdsa.GenerateKey(elliptic.P521(), rand.Reader)

	duration30Days, _ := time.ParseDuration("-30h")
	notBefore := time.Now().Add(duration30Days) // valid 30 hours ago
	duration1Year, _ := time.ParseDuration("90h")
	notAfter := notBefore.Add(duration1Year) // for 90 hours

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, _ := rand.Int(rand.Reader, serialNumberLimit)

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"gIRC-Go Co"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA: true,
	}

	template.IPAddresses = append(template.IPAddresses, net.ParseIP("127.0.0.1"))
	template.IPAddresses = append(template.IPAddresses, net.ParseIP("::"))
	template.DNSNames = append(template.DNSNames, "localhost")

	derBytes, _ := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)

	c := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	b, _ := x509.MarshalECPrivateKey(priv)
	k := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: b})

	// we mock up a server connection to test the client
	listenerKeyPair, _ := tls.X509KeyPair(c, k)

	var listenerTLSConfig tls.Config
	listenerTLSConfig.Certificates = make([]tls.Certificate, 0)
	listenerTLSConfig.Certificates = append(listenerTLSConfig.Certificates, listenerKeyPair)
	listener, _ := tls.Listen("tcp", ":0", &listenerTLSConfig)

	// mock up the client side too
	clientTLSCertPool := x509.NewCertPool()
	clientTLSCertPool.AppendCertsFromPEM(c)

	var clientTLSConfig tls.Config
	clientTLSConfig.RootCAs = clientTLSCertPool
	clientTLSConfig.ServerName = "localhost"
	go client.Connect(listener.Addr().String(), true, &clientTLSConfig)
	go client.ReceiveLoop()

	testServerConnection(t, reactor, client, listener)
}

// marshalPKCS8PrivateKey marshals the provided ECDSA private key into the
// PKCS#8 private key format.
func marshalPKCS8PrivateKey(key *ecdsa.PrivateKey) ([]byte, error) {
	oid, ok := oidFromNamedCurve(key.PublicKey.Curve)
	if !ok {
		return nil, fmt.Errorf("illegal curve")
	}
	paramBytes, err := asn1.Marshal(oid)
	if err != nil {
		return nil, err
	}
	var algo pkix.AlgorithmIdentifier
	algo.Algorithm = oidPublicKeyECDSA
	algo.Parameters.FullBytes = paramBytes

	privBytes, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, err
	}
	pkcs8 := struct {
		Version    int
		Algo       pkix.AlgorithmIdentifier
		PrivateKey []byte
	}{
		Version:    1,
		Algo:       algo,
		PrivateKey: privBytes,
	}
	return asn1.Marshal(pkcs8)
}