// Read a PEM file and ask for a password to decrypt it if needed
func ReadPEMData(pemFile string, pemPass []byte) ([]byte, error) {
	pemData, err := ioutil.ReadFile(pemFile)
	if err != nil {
		return pemData, err
	}

	// We should really just get the pem.Block back here, if there's other
	// junk on the end, warn about it.
	pemBlock, rest := pem.Decode(pemData)
	if len(rest) > 0 {
		log.Warning("Didn't parse all of", pemFile)
	}

	if x509.IsEncryptedPEMBlock(pemBlock) {
		// Decrypt and get the ASN.1 DER bytes here
		pemData, err = x509.DecryptPEMBlock(pemBlock, pemPass)
		if err != nil {
			return pemData, err
		} else {
			log.Info("Decrypted", pemFile, "successfully")
		}
		// Shove the decrypted DER bytes into a new pem Block with blank headers
		var newBlock pem.Block
		newBlock.Type = pemBlock.Type
		newBlock.Bytes = pemData
		// This is now like reading in an uncrypted key from a file and stuffing it
		// into a byte stream
		pemData = pem.EncodeToMemory(&newBlock)
	}
	return pemData, nil
}

// CertificateLockFile adds a new certificate lock on the given Client and
// Config that ensures that a server's certificate is signed by the same CA
// from connection-to-connection. This is helpful when connecting to servers
// with self-signed certificates.
//
// If filename does not exist, the server's certificate chain will be written
// to that file. If it does exist, certificates will be read from that file and
// added to RootCAs in config's TLSConfig.
//
// Example:
//
//  if firstConnectionToServer {
//      // Allow self-signed certificates to be accepted on the initial
//      // connection.
//      config.TLSConfig.InsecureSkipVerify = true
//  }
//  gumbleutil.CertificateLockFile(client, &config, filename)
//
//  if err := client.Connect(); err != nil {
//      panic(err)
//  }
func CertificateLockFile(client *gumble.Client, config *gumble.Config, filename string) (gumble.Detacher, error) {
	if file, err := os.Open(filename); err == nil {
		defer file.Close()
		if config.TLSConfig.RootCAs == nil {
			config.TLSConfig.RootCAs = x509.NewCertPool()
		}
		if data, err := ioutil.ReadAll(file); err == nil {
			config.TLSConfig.RootCAs.AppendCertsFromPEM(data)
		}
		return nil, nil
	}

	return client.Attach(Listener{
		Connect: func(e *gumble.ConnectEvent) {
			tlsClient, ok := e.Client.Conn().(*tls.Conn)
			if !ok {
				return
			}
			serverCerts := tlsClient.ConnectionState().PeerCertificates
			file, err := os.Create(filename)
			if err != nil {
				return
			}
			block := pem.Block{
				Type: "CERTIFICATE",
			}
			for _, cert := range serverCerts {
				block.Bytes = cert.Raw
				pem.Encode(file, &block)
			}
			file.Close()
		},
	}), nil
}

func (pk *PrivateKey) Marshal(ss ...string) ([]byte, error) {
	var k = x509.MarshalPKCS1PrivateKey(&pk.PrivateKey)
	var block pem.Block
	block.Bytes = k
	block.Type = strings.Join(ss, " ")
	return pem.EncodeToMemory(&block), nil
}

func execSSH(url, uuid string, key *rsa.PrivateKey, args []string) error {
	f, err := ioutil.TempFile("", "runx")
	if err != nil {
		return fmt.Errorf("tmpfile: %s", err)
	}
	defer f.Close()
	var b pem.Block
	b.Type = "RSA PRIVATE KEY"
	b.Bytes = x509.MarshalPKCS1PrivateKey(key)
	if err = pem.Encode(f, &b); err != nil {
		return fmt.Errorf("pem: %s", err)
	}
	f.Seek(0, 0)
	argv := []string{
		"ssh",
		"-i" + f.Name(),
		"-oProxyCommand=hk runx [proxy]",
		"-oLocalCommand=rm " + f.Name(),
		"-oStrictHostKeyChecking=no",
		"-oUserKnownHostsFile=/dev/null",
		"[email protected]" + uuid,
	}
	env := append(os.Environ(), "RUNX_URL="+url)
	return syscall.Exec("/usr/bin/ssh", append(argv, args...), env)
}

func EncodePEM(binary []byte, blockType string, password string) (pemBlock string, err error) {

	var blk *pem.Block
	/* Awaiting Go 1.1 */
	if password != "" {
		passwordBytes := ([]byte)(password)
		blk, err = x509.EncryptPEMBlock(rand.Reader, blockType, binary, passwordBytes, x509.PEMCipherAES256)
		if err != nil {
			return
		}
	} else {
		/* */
		blk = new(pem.Block)
		blk.Type = blockType
		blk.Bytes = binary
		/* Awaiting Go 1.1 */
	}
	/* */

	buf := new(bytes.Buffer)

	err = pem.Encode(buf, blk)
	if err != nil {
		return
	}

	pemBlock = buf.String()
	return
}

// ToCSRBundle converts a byte-based raw DER certificate bundle
// to a PEM-based string certificate bundle
func (p *ParsedCSRBundle) ToCSRBundle() (*CSRBundle, error) {
	result := &CSRBundle{}
	block := pem.Block{
		Type: "CERTIFICATE REQUEST",
	}

	if p.CSRBytes != nil && len(p.CSRBytes) > 0 {
		block.Bytes = p.CSRBytes
		result.CSR = strings.TrimSpace(string(pem.EncodeToMemory(&block)))
	}

	if p.PrivateKeyBytes != nil && len(p.PrivateKeyBytes) > 0 {
		block.Bytes = p.PrivateKeyBytes
		switch p.PrivateKeyType {
		case RSAPrivateKey:
			result.PrivateKeyType = "rsa"
			block.Type = "RSA PRIVATE KEY"
		case ECPrivateKey:
			result.PrivateKeyType = "ec"
			block.Type = "EC PRIVATE KEY"
		default:
			return nil, errutil.InternalError{"Could not determine private key type when creating block"}
		}
		result.PrivateKey = strings.TrimSpace(string(pem.EncodeToMemory(&block)))
	}

	return result, nil
}

func GenKeyPairIfNone(privateName string, publicName string) {
	dir, _ := filepath.Abs(filepath.Dir(os.Args[0]))
	privatekey := filepath.Join(dir, privateName)
	publickey := filepath.Join(dir, publicName)

	if _, err := os.Stat(string(privatekey)); os.IsNotExist(err) {

		log.Println("Generating JWT private key at ", string(privatekey))
		k, err := rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			log.Fatal(err)
		}
		var private pem.Block
		private.Type = "RSA PRIVATE KEY"
		private.Bytes = x509.MarshalPKCS1PrivateKey(k)
		pp := new(bytes.Buffer)
		pem.Encode(pp, &private)
		err = ioutil.WriteFile(string(privatekey), pp.Bytes(), 0644)
		if err != nil {
			log.Fatal(err)
		}

		log.Println("Generating JWT public key at ", string(privatekey))
		var public pem.Block
		public.Type = "RSA PUBLIC KEY"
		public.Bytes, _ = x509.MarshalPKIXPublicKey(&k.PublicKey)
		ps := new(bytes.Buffer)
		pem.Encode(ps, &public)
		err = ioutil.WriteFile(string(publickey), ps.Bytes(), 0644)
		if err != nil {
			log.Fatal(err)
		}
	}

}

// Given a private key and a (possibly empty) password, returns a byte
// slice containing a PEM-encoded private key in the appropriate
// OpenSSH format.
func MarshalPrivate(priv interface{}, password string) (out []byte, err error) {
	var (
		keytype Type
		der     []byte
		btype   string
	)

	switch priv.(type) {
	case *rsa.PrivateKey:
		keytype = KEY_RSA
		der = x509.MarshalPKCS1PrivateKey(priv.(*rsa.PrivateKey))
		if der == nil {
			err = ErrInvalidPrivateKey
			return
		}
		btype = "RSA PRIVATE KEY"
	case *ecdsa.PrivateKey:
		keytype = KEY_ECDSA
		der, err = marshalECDSAKey(priv.(*ecdsa.PrivateKey))
		btype = "EC PRIVATE KEY"
	case *dsa.PrivateKey:
		keytype = KEY_DSA

		dsakey := priv.(*dsa.PrivateKey)
		k := struct {
			Version int
			P       *big.Int
			Q       *big.Int
			G       *big.Int
			Priv    *big.Int
			Pub     *big.Int
		}{
			Version: 1,
			P:       dsakey.PublicKey.P,
			Q:       dsakey.PublicKey.Q,
			G:       dsakey.PublicKey.G,
			Priv:    dsakey.PublicKey.Y,
			Pub:     dsakey.X,
		}
		der, err = asn1.Marshal(k)
		if err != nil {
			return
		}
		btype = "DSA PRIVATE KEY"
	default:
		err = ErrInvalidPrivateKey
		return
	}

	if password != "" {
		out, err = encrypt(der, keytype, password)
		return
	}
	var block pem.Block
	block.Type = btype
	block.Bytes = der
	out = pem.EncodeToMemory(&block)
	return
}

func readPEMFile(path, passphrase string) ([]byte, error) {
	pass := []byte(passphrase)
	var blocks []*pem.Block

	content, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}

	for len(content) > 0 {
		var block *pem.Block

		block, content = pem.Decode(content)
		if block == nil {
			if len(blocks) == 0 {
				return nil, errors.New("no pem file")
			}
			break
		}

		if x509.IsEncryptedPEMBlock(block) {
			var buffer []byte
			var err error
			if len(pass) == 0 {
				err = errors.New("No passphrase available")
			} else {
				// Note, decrypting pem might succeed even with wrong password, but
				// only noise will be stored in buffer in this case.
				buffer, err = x509.DecryptPEMBlock(block, pass)
			}

			if err != nil {
				logp.Err("Dropping encrypted pem '%v' block read from %v. %v",
					block.Type, path, err)
				continue
			}

			// DEK-Info contains encryption info. Remove header to mark block as
			// unencrypted.
			delete(block.Headers, "DEK-Info")
			block.Bytes = buffer
		}
		blocks = append(blocks, block)
	}

	if len(blocks) == 0 {
		return nil, errors.New("no PEM blocks")
	}

	// re-encode available, decrypted blocks
	buffer := bytes.NewBuffer(nil)
	for _, block := range blocks {
		err := pem.Encode(buffer, block)
		if err != nil {
			return nil, err
		}
	}
	return buffer.Bytes(), nil
}

func byteToPEM(b []byte, t PEMType) []byte {
	var blk pem.Block = pem.Block{
		Type:    string(t),
		Headers: nil,
	}

	blk.Bytes = b
	data := pem.EncodeToMemory(&blk)
	return data
}

func main() {
	k, err := rsa.GenerateKey(rand.Reader, 768)
	if err != nil {
		log.Fatal(err)
	}
	var b pem.Block
	b.Type = "RSA PRIVATE KEY"
	b.Bytes = x509.MarshalPKCS1PrivateKey(k)
	pem.Encode(os.Stdout, &b)
}

func PrintPublicKey(pubkey crypto.PublicKey) {

	bytes, _ := x509.MarshalPKIXPublicKey(pubkey)
	block := pem.Block{}
	block.Type = "EC PUBLIC KEY"

	block.Bytes = bytes
	bytes_encoded := pem.EncodeToMemory(&block)

	log.Debugf("Public key:\n%s", string(bytes_encoded))
}

// ExportPrivatePEM writes the RSA private key to a file in PEM format.
func ExportPrivatePEM(prv *rsa.PrivateKey, filename string) (err error) {
	cert := x509.MarshalPKCS1PrivateKey(prv)
	blk := new(pem.Block)
	blk.Type = "RSA PRIVATE KEY"
	blk.Bytes = cert
	out, err := os.Create(filename)
	if err == nil {
		err = pem.Encode(out, blk)
	}
	return
}

func main() {
	var blk pem.Block = pem.Block{
		Type:    "CERTIFICATE",
		Headers: nil,
	}

	ca_b, _ := ioutil.ReadFile("ca.pem")
	blk.Bytes = ca_b

	data := pem.EncodeToMemory(&blk)
	fmt.Printf("%s\n", data)
}

// ExportPublicPEM writes the public key to a file in PEM format.
func ExportPublicPEM(pub *rsa.PublicKey, filename string) (err error) {
	cert, err := x509.MarshalPKIXPublicKey(pub)
	if err != nil {
		return
	}
	blk := new(pem.Block)
	blk.Type = "RSA PUBLIC KEY"
	blk.Bytes = cert
	out, err := os.Create(filename)
	if err == nil {
		err = pem.Encode(out, blk)
	}
	return
}

func newECDSACertificateRequest(priv *ecdsa.PrivateKey, si *SubjectInfo) (out []byte, err error) {
	var siAttr rdnSequence

	rdnAppendPrintable(si.Country, asn1CountryName, &siAttr)
	rdnAppendPrintable(si.StateOrProvince, asn1StateOrProvName, &siAttr)
	rdnAppendPrintable(si.Locality, asn1LocalityName, &siAttr)
	rdnAppendPrintable(si.OrgName, asn1OrgName, &siAttr)
	rdnAppendPrintable(si.OrgUnitName, asn1OrgUnitName, &siAttr)
	rdnAppendPrintable(si.CommonName, asn1CommonName, &siAttr)
	rdnAppendPrintable(si.Email, asn1EmailAddress, &siAttr)

	pkInfo, err := encodeECDSA(priv.PublicKey)
	if err != nil {
		return
	}
	var csrInfo = certificateRequestInfo{
		Subject: siAttr,
		PKInfo:  pkInfo,
	}

	sigData, err := asn1.Marshal(csrInfo)
	if err != nil {
		return
	}
	sum := sha256.Sum256(sigData)
	r, s, err := ecdsa.Sign(rand.Reader, priv, sum[:])
	if err != nil {
		return
	}
	ecdsaSig := ecdsaSignature{r, s}
	sig, err := asn1.Marshal(ecdsaSig)
	if err != nil {
		return
	}

	var crt = certificateRequest{
		Info:      csrInfo,
		SigAlgo:   nullAlgorithm(asn1SHA256withECDSA),
		Signature: toBitString(sig),
	}

	var block pem.Block
	block.Type = "CERTIFICATE REQUEST"
	block.Bytes, err = asn1.Marshal(crt)
	if err != nil {
		return
	}
	out = pem.EncodeToMemory(&block)
	return
}

func generatePrivatePem(prv *rsa.PrivateKey) (prvPem string, err error) {
	cert := x509.MarshalPKCS1PrivateKey(prv)
	blk := new(pem.Block)
	blk.Type = "RSA PRIVATE KEY"
	blk.Bytes = cert

	var b bytes.Buffer
	err = pem.Encode(&b, blk)
	if err != nil {
		return
	}

	prvPem = b.String()
	return
}

// packageBox actually handles boxing. It can output either PEM-encoded or
// DER-encoded boxes.
func packageBox(lockedKey, box []byte, armour bool) (pkg []byte, err error) {
	var pkgBox = boxPackage{lockedKey, box}

	pkg, err = asn1.Marshal(pkgBox)
	if err != nil {
		fmt.Println("[!] couldn't package the box")
		return
	}

	if armour {
		var block pem.Block
		block.Type = "SSHBOX ENCRYPTED FILE"
		block.Bytes = pkg
		pkg = pem.EncodeToMemory(&block)
	}
	return
}

// GenerateKeyAndCSR generates a new PEM-encoded private key and
// certificate signature request. The keyType parameter controls what
// type of key is generated; the size parameter controls what size key
// (in bits) is generated. For ECDSA, valid key sizes are 256, 384,
// and 521; these translate to secp256v1, secp384r1, and secp521r1. The
// SubjectInfo parameter contains the information used to create the CSR.
func GenerateKeyAndCSR(keyType PublicAlgo, size int, si *SubjectInfo) (pemKey []byte, csr []byte, err error) {
	var priv interface{}
	var block pem.Block

	switch keyType {
	case RSA:
		priv, err = rsa.GenerateKey(rand.Reader, size)
		if err != nil {
			return nil, nil, err
		}
		block.Bytes = x509.MarshalPKCS1PrivateKey(priv.(*rsa.PrivateKey))
		block.Type = "RSA PRIVATE KEY"
	case ECDSA:
		var curve elliptic.Curve
		switch size {
		case 256:
			curve = elliptic.P256()
		case 384:
			curve = elliptic.P384()
		case 521:
			curve = elliptic.P521()
		default:
			err = ErrUnsupportedPrivate
			return
		}

		priv, err = ecdsa.GenerateKey(curve, rand.Reader)
		if err != nil {
			return
		}

		block.Bytes, err = x509.MarshalECPrivateKey(priv.(*ecdsa.PrivateKey))
		if err != nil {
			return
		}
		block.Type = "EC PRIVATE KEY"
	default:
		err = ErrUnsupportedPrivate
		return
	}

	pemKey = pem.EncodeToMemory(&block)
	csr, err = NewCertificateRequest(priv, si)
	return
}

// CertificateLockFile adds a new certificate lock on the given Client that
// ensures that a server's certificate chain is the same from
// connection-to-connection. This is helpful when connecting to servers with
// self-signed certificates.
//
// If filename does not exist, the server's certificate chain will be written
// to that file. If it does exist, certificates will be read from the file and
// checked against the server's certificate chain upon connection.
//
// Example:
//
//  if allowSelfSignedCertificates {
//      config.TLSConfig.InsecureSkipVerify = true
//  }
//  gumbleutil.CertificateLockFile(client, filename)
//
//  if err := client.Connect(); err != nil {
//      panic(err)
//  }
func CertificateLockFile(client *gumble.Client, filename string) {
	client.Config.TLSVerify = func(state *tls.ConnectionState) error {
		if file, err := os.Open(filename); err == nil {
			defer file.Close()
			data, err := ioutil.ReadAll(file)
			if err != nil {
				return err
			}
			i := 0
			for block, data := pem.Decode(data); block != nil; block, data = pem.Decode(data) {
				cert, err := x509.ParseCertificate(block.Bytes)
				if err != nil {
					return err
				}
				if i >= len(state.PeerCertificates) {
					return errors.New("gumbleutil: invalid certificate chain length")
				}
				if !cert.Equal(state.PeerCertificates[i]) {
					return errors.New("gumbleutil: certificate verification failure")
				}
				i++
			}
			if i != len(state.PeerCertificates) {
				return errors.New("gumbleutil: invalid certificate chain length")
			}
			return nil
		}

		file, err := os.Create(filename)
		if err != nil {
			return err
		}
		defer file.Close()
		block := pem.Block{
			Type: "CERTIFICATE",
		}
		for _, cert := range state.PeerCertificates {
			block.Bytes = cert.Raw
			if err := pem.Encode(file, &block); err != nil {
				return err
			}
		}
		return nil
	}
}