package encryptor

import (
	"bytes"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path"
	"strings"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/packet"
	// for openpgp
	_ "golang.org/x/crypto/ripemd160"
)

const keylen = 32

// Encryptor interface
type Encryptor interface {
	Encrypt(string) string
	Decrypt(string) string
	SetPublic(string)
	GetPublic() string
	SetSymmetric(string)

	FromFiles(string, string) error
	ToFiles(string, string) error
}

type encryptor struct {
	entity    *openpgp.Entity
	writing   *openpgp.Entity
	symmetric string
	b64       bool
}

func (e *encryptor) SetPublic(pub string) {
	reader := b64dec(pub)
	e.writing = entityFromPackets(nil, strToPubPack(reader))
}

func (e *encryptor) GetPublic() string {
	bufpub := new(bytes.Buffer)
	e.writing.Serialize(bufpub)
	bufPu, err := ioutil.ReadAll(bufpub)
	if err != nil {
		panic(err)
	}
	return b64enc(string(bufPu))
}

func (e *encryptor) setPrivate(pri string) {
	readerPri := b64dec(pri)
	readerPub := b64dec(e.GetPublic())
	e.entity = entityFromPackets(strToPriPack(readerPri), strToPubPack(readerPub))
}

func (e *encryptor) getPrivate() string {
	bufpri := new(bytes.Buffer)
	e.entity.SerializePrivate(bufpri, nil)
	bufPr, err := ioutil.ReadAll(bufpri)
	if err != nil {
		panic(err)
	}
	return b64enc(string(bufPr))
}

// NewKeyPair generates a PGP private-public key pair.
func NewKeyPair() (string, string) {
	entity, err := openpgp.NewEntity("", "", "", nil)
	if err != nil {
		panic(err)
	}
	en := &encryptor{
		entity:  entity,
		writing: entity,
	}
	return en.getPrivate(), en.GetPublic()
}

func (e *encryptor) getPubPacket() *packet.PublicKey {
	return e.writing.PrimaryKey
}

func (e *encryptor) getPriPacket() *packet.PrivateKey {
	return e.entity.PrivateKey
}

func entityFromPackets(priKey *packet.PrivateKey, pubKey *packet.PublicKey) *openpgp.Entity {
	config := packet.Config{
		DefaultHash:            crypto.SHA256,
		DefaultCipher:          packet.CipherAES256,
		DefaultCompressionAlgo: packet.CompressionZLIB,
		CompressionConfig: &packet.CompressionConfig{
			Level: 9,
		},
		RSABits: 4096,
	}

	currentTime := config.Now()
	uid := packet.NewUserId("", "", "")

	e := openpgp.Entity{
		PrimaryKey: pubKey,
		PrivateKey: priKey,
		Identities: make(map[string]*openpgp.Identity),
	}
	isPrimaryID := false

	e.Identities[uid.Id] = &openpgp.Identity{
		Name:   uid.Name,
		UserId: uid,
		SelfSignature: &packet.Signature{
			CreationTime: currentTime,
			SigType:      packet.SigTypePositiveCert,
			PubKeyAlgo:   packet.PubKeyAlgoRSA,
			Hash:         config.Hash(),
			IsPrimaryId:  &isPrimaryID,
			FlagsValid:   true,
			FlagSign:     true,
			FlagCertify:  true,
			IssuerKeyId:  &e.PrimaryKey.KeyId,
		},
	}

	keyLifetimeSecs := uint32(86400 * 365)

	e.Subkeys = make([]openpgp.Subkey, 1)
	e.Subkeys[0] = openpgp.Subkey{
		PublicKey:  pubKey,
		PrivateKey: priKey,
		Sig: &packet.Signature{
			CreationTime:              currentTime,
			SigType:                   packet.SigTypeSubkeyBinding,
			PubKeyAlgo:                packet.PubKeyAlgoRSA,
			Hash:                      config.Hash(),
			PreferredHash:             []uint8{8}, // SHA-256
			FlagsValid:                true,
			FlagEncryptStorage:        true,
			FlagEncryptCommunications: true,
			IssuerKeyId:               &e.PrimaryKey.KeyId,
			KeyLifetimeSecs:           &keyLifetimeSecs,
		},
	}
	return &e

}

func strToPack(str string) interface{} {
	reader := b64dec(str)
	KeyRead := bytes.NewReader([]byte(reader))
	PacketReader := packet.NewReader(KeyRead)
	pkt, err := PacketReader.Next()
	if err != nil {
		panic(err)
	}
	return pkt
}

func strToPriPack(private string) *packet.PrivateKey {
	priKeyPack, ok := strToPack(private).(*packet.PrivateKey)
	if !ok {
		panic(ok)
	}
	return priKeyPack
}

func strToPubPack(public string) *packet.PublicKey {
	pubKeyPack, ok := strToPack(public).(*packet.PublicKey)
	if !ok {
		panic(ok)
	}
	return pubKeyPack
}

// NewEncryptor function
func NewEncryptor(private, public string, b64 ...bool) *encryptor {
	backupPri, backupPub := NewKeyPair()
	if private == "" {
		private = backupPri
	}
	if public == "" {
		public = backupPub
	}
	priKeyPack := strToPriPack(private)
	pubKeyPack := strToPubPack(public)

	b64payloads := len(b64) > 0 && b64[0]

	return &encryptor{
		entity:  entityFromPackets(priKeyPack, pubKeyPack),
		writing: entityFromPackets(nil, pubKeyPack),
		b64:     b64payloads,
	}
}

func b64enc(msg string) string {
	return base64.StdEncoding.EncodeToString([]byte(msg))
}

func b64dec(msg string) string {
	decMe, err := base64.StdEncoding.DecodeString(msg)
	if err != nil {
		return msg
	}
	return string(decMe)
}

func (e *encryptor) SetSymmetric(key string) {
	e.symmetric = key
}

func (e *encryptor) Encrypt(msg string) string {
	if e.symmetric != "" {
		return e.symEncrypt(msg)
	}
	return e.asymEncrypt(msg)
}

func (e *encryptor) Decrypt(msg string) string {
	if e.symmetric != "" {
		return e.symDecrypt(msg)
	}
	return e.asymDecrypt(msg)
}

func fixKey(key string) []byte {
	return []byte(strings.Repeat("minkey"+string(key), keylen)[:keylen])
}

func pad(src []byte) []byte {
	padding := aes.BlockSize - len(src)%aes.BlockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}

func unpad(src []byte) ([]byte, error) {
	length := len(src)
	unpadding := int(src[length-1])

	if unpadding > length {
		return nil, errors.New("unpad error. This could happen when incorrect encryption key is used")
	}

	return src[:(length - unpadding)], nil
}

func (e *encryptor) symEncrypt(smsg string) string {
	key := fixKey(e.symmetric)
	block, err := aes.NewCipher(key)
	if err != nil {
		return ""
	}

	msg := pad([]byte(smsg))
	ciphertext := make([]byte, aes.BlockSize+len(msg))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return ""
	}

	cfb := cipher.NewCFBEncrypter(block, iv)
	cfb.XORKeyStream(ciphertext[aes.BlockSize:], msg)
	if e.b64 {
		return b64enc(string(ciphertext))
	}
	return string(ciphertext)
}

func (e *encryptor) symDecrypt(msg string) string {
	payload := []byte(msg)
	if e.b64 {
		payload = []byte(b64dec(msg))
	}
	if len(payload) == 0 {
		return ""
	}
	key := fixKey(e.symmetric)
	block, err := aes.NewCipher(key)
	if err != nil {
		return ""
	}

	if (len(payload) % aes.BlockSize) != 0 {
		return ""
	}

	iv := payload[:aes.BlockSize]
	bmsg := payload[aes.BlockSize:]

	cfb := cipher.NewCFBDecrypter(block, iv)
	cfb.XORKeyStream(bmsg, bmsg)

	b, _ := unpad(bmsg)
	return string(b)
}

// Encrypt function
func (e *encryptor) asymEncrypt(msg string) string {
	// encrypt string
	buf := new(bytes.Buffer)
	w, err := openpgp.Encrypt(buf, []*openpgp.Entity{e.writing}, nil, nil, nil)
	if err != nil {
		return ""
	}
	_, err = w.Write([]byte(msg))
	if err != nil {
		return ""
	}
	err = w.Close()
	if err != nil {
		return ""
	}
	// Encode to b64
	return b64enc(string(buf.Bytes()))
}

// Decrypt function
func (e *encryptor) asymDecrypt(msg string) string {
	// Decode the b64 string
	decMe := b64dec(msg)
	KeyRead := bytes.NewReader([]byte(decMe))

	md, err := openpgp.ReadMessage(KeyRead, openpgp.EntityList([]*openpgp.Entity{e.entity}), nil, nil)
	if err != nil {
		return ""
	}

	plaintext, err := ioutil.ReadAll(md.UnverifiedBody)
	if err != nil {
		return ""
	}

	return string(plaintext)
}

func (e *encryptor) FromFiles(privatePath, publicPath string) error {
	priRaw, err := ioutil.ReadFile(privatePath)
	if err != nil {
		return err
	}
	pubRaw, err := ioutil.ReadFile(publicPath)
	if err != nil {
		return err
	}

	d := NewEncryptor(string(priRaw), string(pubRaw))

	v := d.Encrypt("Hello")
	if len(v) == 0 {
		return errors.New("bad public-private key file pair")
	}
	v = d.Decrypt(v)
	if len(v) == 0 {
		return errors.New("bad public-private key file pair")
	}
	e.entity = d.entity
	e.writing = d.entity
	return nil
}

func (e *encryptor) ToFiles(privatePath, publicPath string) error {
	d := &encryptor{
		writing: e.entity,
		entity:  e.entity,
	}
	if err := os.MkdirAll(path.Dir(privatePath), os.ModePerm); err != nil {
		return err
	}
	if err := ioutil.WriteFile(privatePath, []byte(d.getPrivate()), os.ModePerm); err != nil {
		return err
	}
	if err := os.MkdirAll(path.Dir(publicPath), os.ModePerm); err != nil {
		return err
	}
	return ioutil.WriteFile(publicPath, []byte(d.GetPublic()), os.ModePerm)
}