// Copyright 2022 Martin Dosch.
// Use of this source code is governed by the BSD-2-clause
// license that can be found in the LICENSE file.

package main

import (
	"encoding/base64"
	"errors"
	"log"
	"os"
	"runtime"
	"strings"
	"time"

	"github.com/ProtonMail/gopenpgp/v2/crypto" // MIT License
	"github.com/beevik/etree"                  // BSD-2-clause
	"github.com/mattn/go-xmpp"                 // BSD-3-Clause
)

func oxDeleteNodes(jid string, client *xmpp.Client,
	iqc chan xmpp.IQ) error {
	nlr, err := oxNodeListRequest()
	if err != nil {
		return err
	}
	iqReply, err := sendIQ(client, iqc, jid, "get", nlr)
	if err != nil {
		return err
	}
	nodes, err := oxNodeListReply(iqReply.Query)
	if err != nil {
		return err
	}
	for _, node := range nodes {
		dnr, err := oxDeleteNodeRequest(node)
		if err != nil {
			continue
		}
		_, err = sendIQ(client, iqc, jid, "set", dnr)
		if err != nil {
			continue
		}
	}
	return nil
}

func oxDecrypt(m xmpp.Chat, client *xmpp.Client, iqc chan xmpp.IQ,
	user string, oxPrivKey *crypto.Key) (string, time.Time, error) {
	var cryptMsgByte []byte
	var err error
	sender := strings.Split(m.Remote, "/")[0]
	for _, r := range m.OtherElem {
		if r.XMLName.Space == nsOx {
			cryptMsgByte, err =
				base64.StdEncoding.DecodeString(r.InnerXML)
			if err != nil {
				return "error", time.Now(), err
			}
			break
		}
	}
	oxMsg := crypto.NewPGPMessage(cryptMsgByte)
	keyRing, err := crypto.NewKeyRing(oxPrivKey)
	if err != nil {
		return "error", time.Now(), err
	}
	senderKeyRing, err := oxGetPublicKeyRing(client, iqc, sender)
	if err != nil {
		return "error", time.Now(), err
	}
	decryptMsg, err := keyRing.Decrypt(oxMsg, senderKeyRing, crypto.GetUnixTime())
	if err != nil {
		return "error", time.Now(), err
	}
	// Remove invalid code points.
	message := validUTF8(string(decryptMsg.Data))
	signcrypt, err := oxParseSignCrypt(message)
	if err != nil {
		return "error", time.Now(), err
	}
	if strings.Split(signcrypt.body, "/")[0] != user {
		return "error", time.Now(), errors.New("ox: encrypted for wrong user")
	}
	msgStamp, err := time.Parse("2006-01-02T15:04:05Z0700", signcrypt.stamp)
	if err != nil {
		return "error", time.Now(), err
	}
	return signcrypt.body, msgStamp, nil
}

func isOxMsg(m xmpp.Chat) bool {
	for _, r := range m.OtherElem {
		if r.XMLName.Space == nsOx {
			return true
		}
	}
	return false
}

func oxImportPrivKey(jid string, privKeyLocation string, client *xmpp.Client,
	iqc chan xmpp.IQ) error {
	xmppURI := "xmpp:" + jid
	buffer, err := readFile(privKeyLocation)
	if err != nil {
		return err
	}
	key, err := crypto.NewKey(buffer.Bytes())
	if err != nil {
		key, err = crypto.NewKeyFromArmored(buffer.String())
		if err != nil {
			keyDecoded, err := base64.StdEncoding.DecodeString(buffer.String())
			if err != nil {
				return err
			}
			key, err = crypto.NewKey(keyDecoded)
			if err != nil {
				return err
			}
		}
	}
	entity := key.GetEntity()
	if entity.Identities[xmppURI] == nil {
		return errors.New("Key identity is not " + xmppURI)
	}
	pk, err := key.GetPublicKey()
	if err != nil {
		return err
	}
	pubKey, err := crypto.NewKey(pk)
	if err != nil {
		return err
	}
	fingerprint := strings.ToUpper(pubKey.GetFingerprint())
	_, err = oxRecvPublicKeys(client, iqc, jid, fingerprint)
	if err != nil {
		err = oxPublishPubKey(jid, client, iqc, pubKey)
		if err != nil {
			return err
		}
	}
	location, err := oxGetPrivKeyLoc(jid)
	if err != nil {
		return err
	}
	keySerialized, err := key.Serialize()
	if err != nil {
		return err
	}
	err = oxStoreKey(location,
		base64.StdEncoding.EncodeToString(keySerialized))
	if err != nil {
		log.Fatal(err)
	}
	pubKeyRing, err := oxGetPublicKeyRing(client, iqc, jid)
	if err == nil {
		pubKeys := pubKeyRing.GetKeys()
		for _, r := range pubKeys {
			if strings.ToUpper(r.GetFingerprint()) == fingerprint {
				return nil
			}
		}
	}
	err = oxPublishPubKey(jid, client, iqc, pubKey)
	if err != nil {
		return err
	}
	return nil
}

func oxPublishPubKey(jid string, client *xmpp.Client, iqc chan xmpp.IQ,
	pubKey *crypto.Key) error {
	keyCreated := time.Now().UTC().Format("2006-01-02T15:04:05Z")
	fingerprint := strings.ToUpper(pubKey.GetFingerprint())
	keySerialized, err := pubKey.Serialize()
	if err != nil {
		return err
	}
	pubKeyBase64 := base64.StdEncoding.EncodeToString(keySerialized)
	ppk, err := oxPublishPubKeyRequest(pubKeyBase64, fingerprint, keyCreated)
	if err != nil {
		return err
	}
	iqReply, err := sendIQ(client, iqc, jid, "set", ppk)
	if err != nil {
		return err
	}
	if iqReply.Type != "result" {
		return errors.New("error while publishing public key")
	}
	ownPubKeyRingFromPubsub, err := oxRecvPublicKeys(client, iqc, jid, fingerprint)
	if err != nil {
		return errors.New("couldn't successfully verify public key upload")
	}
	ownPubKeyFromPubsub := ownPubKeyRingFromPubsub.GetKeys()[0]
	ownPubKeyFromPubsubSerialized, err := ownPubKeyFromPubsub.Serialize()
	if err != nil {
		return errors.New("couldn't successfully verify public key upload")
	}
	if pubKeyBase64 != base64.StdEncoding.EncodeToString(ownPubKeyFromPubsubSerialized) {
		return errors.New("couldn't successfully verify public key upload")
	}
	ppkl, err := oxPublishPubKeyListRequest(fingerprint, keyCreated)
	if err != nil {
		return err
	}
	iqReply, err = sendIQ(client, iqc, jid, "set", ppkl)
	if err != nil {
		return err
	}
	if iqReply.Type != "result" {
		return errors.New("couldn't publish public key list")
	}
	return nil
}

func oxGetPrivKeyLoc(jid string) (string, error) {
	var err error
	var homeDir, dataDir string
	switch {
	case os.Getenv("$XDG_DATA_HOME") != "":
		dataDir = os.Getenv("$XDG_DATA_HOME")
	case os.Getenv("$XDG_HOME") != "":
		homeDir = os.Getenv("$XDG_HOME")
		dataDir = homeDir + "/.local/share"
	case os.Getenv("$HOME") != "":
		homeDir = os.Getenv("$HOME")
		dataDir = homeDir + "/.local/share"
	default:
		homeDir, err = os.UserHomeDir()
		if err != nil {
			return "error", err
		}
		if homeDir == "" {
			return "error", err
		}
		dataDir = homeDir + "/.local/share"
	}
	dataDir = dataDir + "/go-sendxmpp/oxprivkeys/"
	if _, err = os.Stat(dataDir); os.IsNotExist(err) {
		err = os.MkdirAll(dataDir, 0700)
		if err != nil {
			return "error", err
		}
	}
	dataFile := dataDir + base64.StdEncoding.EncodeToString([]byte(jid))
	return dataFile, nil
}

func oxGetPubKeyLoc(fingerprint string) (string, error) {
	var err error
	var homeDir, dataDir string
	switch {
	case os.Getenv("$XDG_DATA_HOME") != "":
		dataDir = os.Getenv("$XDG_DATA_HOME")
	case os.Getenv("$XDG_HOME") != "":
		homeDir = os.Getenv("$XDG_HOME")
		dataDir = homeDir + "/.local/share"
	case os.Getenv("$HOME") != "":
		homeDir = os.Getenv("$HOME")
		dataDir = homeDir + "/.local/share"
	default:
		homeDir, err = os.UserHomeDir()
		if err != nil {
			return "error", err
		}
		if homeDir == "" {
			return "error", err
		}
		dataDir = homeDir + "/.local/share"
	}
	dataDir = dataDir + "/go-sendxmpp/oxpubkeys/"
	if _, err = os.Stat(dataDir); os.IsNotExist(err) {
		err = os.MkdirAll(dataDir, 0700)
		if err != nil {
			return "error", err
		}
	}
	dataFile := dataDir + fingerprint
	return dataFile, nil
}

func oxGetPrivKey(jid string, passphrase string) (*crypto.Key, error) {
	dataFile, err := oxGetPrivKeyLoc(jid)
	if err != nil {
		log.Fatal(err)
	}
	keyBuffer, err := readFile(dataFile)
	if err != nil {
		log.Fatal(err)
	}
	keyString := keyBuffer.String()
	decodedPrivKey, err := base64.StdEncoding.DecodeString(keyString)
	if err != nil {
		return nil, err
	}
	key, err := crypto.NewKey(decodedPrivKey)
	if err != nil {
		return nil, err
	}
	if passphrase != "" {
		key, err = key.Unlock([]byte(passphrase))
		if err != nil {
			log.Fatal("Ox: couldn't unlock private key.")
		}
	}
	isLocked, err := key.IsLocked()
	if err != nil {
		return nil, err
	}
	if isLocked {
		log.Fatal("Ox: private key is locked.")
	}
	if key.IsExpired() {
		return nil, errors.New("Ox: private key is expired: " + key.GetFingerprint())
	}
	return key, nil
}

func oxStoreKey(location string, key string) error {
	var file *os.File
	file, err := os.Create(location)
	if err != nil {
		return err
	}
	defer file.Close()
	if runtime.GOOS != "windows" {
		_ = file.Chmod(os.FileMode(0600))
	} else {
		_ = file.Chmod(os.FileMode(0200))
	}
	_, err = file.Write([]byte(key))
	if err != nil {
		return err
	}
	return nil
}

func oxGenPrivKey(jid string, client *xmpp.Client, iqc chan xmpp.IQ,
	passphrase string, keyType string) error {
	xmppURI := "xmpp:" + jid
	key, err := crypto.GenerateKey(xmppURI, "", keyType, 4096)
	if err != nil {
		return err
	}
	if passphrase != "" {
		key, err = key.Lock([]byte(passphrase))
		if err != nil {
			return err
		}
	}
	keySerialized, err := key.Serialize()
	if err != nil {
		return err
	}
	location, err := oxGetPrivKeyLoc(jid)
	if err != nil {
		return err
	}
	err = oxStoreKey(location,
		base64.StdEncoding.EncodeToString(keySerialized))
	if err != nil {
		log.Fatal(err)
	}
	decodedPubKey, err := key.GetPublicKey()
	if err != nil {
		return err
	}
	pubKey, err := crypto.NewKey(decodedPubKey)
	if err != nil {
		return err
	}
	err = oxPublishPubKey(jid, client, iqc, pubKey)
	if err != nil {
		return err
	}
	return nil
}

func oxRecvPublicKeys(client *xmpp.Client, iqc chan xmpp.IQ, recipient string,
	fingerprint string) (*crypto.KeyRing, error) {
	opkr, err := oxRecvPubKeysRequest(fingerprint)
	if err != nil {
		return nil, err
	}
	oxPublicKey, err := sendIQ(client, iqc, recipient, "get", opkr)
	if err != nil {
		return nil, err
	}
	if oxPublicKey.Type != "result" {
		return nil, errors.New("Error while requesting public key for " +
			recipient)
	}
	oxPublicKeyXML := etree.NewDocument()
	err = oxPublicKeyXML.ReadFromBytes(oxPublicKey.Query)
	if err != nil {
		return nil, err
	}
	keyring, err := crypto.NewKeyRing(nil)
	if err != nil {
		return nil, err
	}
	oxPublicKeyXMLPubsub := oxPublicKeyXML.SelectElement("pubsub")
	if oxPublicKeyXMLPubsub == nil {
		return nil, errors.New("ox: no pubsub element in reply to public " +
			"key request")
	}
	oxPublicKeyXMLItems := oxPublicKeyXMLPubsub.SelectElement("items")
	if oxPublicKeyXMLItems == nil {
		return nil, errors.New("ox: no items element in reply to public " +
			"key request")
	}
	oxPublicKeyXMLItem := oxPublicKeyXMLItems.SelectElement("item")
	if oxPublicKeyXMLItem == nil {
		return nil, errors.New("ox: no item element in reply to public " +
			"key request")
	}
	oxPublicKeyXMLPubkeys := oxPublicKeyXMLItem.SelectElements("pubkey")
	for _, r := range oxPublicKeyXMLPubkeys {
		data := r.SelectElement("data")
		if data == nil {
			continue
		}
		decodedPubKey, err := base64.StdEncoding.DecodeString(data.Text())
		if err != nil {
			return nil, err
		}
		key, err := crypto.NewKey(decodedPubKey)
		if err != nil {
			return nil, err
		}
		if key.IsExpired() {
			return nil, errors.New("Key is expired: " + fingerprint)
		}
		err = keyring.AddKey(key)
		if err != nil {
			return nil, err
		}
	}
	return keyring, nil
}

func oxGetPublicKeyRing(client *xmpp.Client, iqc chan xmpp.IQ,
	recipient string) (*crypto.KeyRing, error) {
	publicKeyRing, err := crypto.NewKeyRing(nil)
	if err != nil {
		return nil, err
	}
	opkl, err := oxGetPubKeyRingRequest()
	if err != nil {
		return nil, err
	}
	oxPublicKeyList, err := sendIQ(client, iqc, recipient, "get", opkl)
	if err != nil {
		log.Fatal(err)
	}
	if oxPublicKeyList.Type != "result" {
		return nil, errors.New("Error while requesting public openpgp keys for " +
			recipient)
	}
	oxPubKeyListXML := etree.NewDocument()
	err = oxPubKeyListXML.ReadFromBytes(oxPublicKeyList.Query)
	if err != nil {
		return nil, err
	}

	pubKeyRingID := "none"
	newestKey, err := time.Parse(time.RFC3339, "1900-01-01T00:00:00Z")
	if err != nil {
		return nil, err
	}

	oxPubKeyListXMLPubsub := oxPubKeyListXML.SelectElement("pubsub")
	if oxPubKeyListXMLPubsub == nil {
		return nil, errors.New("ox: no pubsub element in public key list")
	}
	oxPubKeyListXMLPubsubItems := oxPubKeyListXMLPubsub.SelectElement("items")
	if oxPubKeyListXMLPubsubItems == nil {
		return nil, errors.New("ox: no items element in public key list")
	}
	oxPubKeyListXMLPubsubItemsItem := oxPubKeyListXMLPubsubItems.SelectElement("item")
	if oxPubKeyListXMLPubsubItemsItem == nil {
		return nil, errors.New("ox: no item element in public key list")
	}
	oxPubKeyListXMLPubsubItemsItemPkl := oxPubKeyListXMLPubsubItemsItem.SelectElement("public-keys-list")
	if oxPubKeyListXMLPubsubItemsItemPkl == nil {
		return nil, errors.New("ox: no keblic-keys-list element")
	}
	oxPubKeyListXMLPubsubItemsItemPklPm :=
		oxPubKeyListXMLPubsubItemsItemPkl.SelectElements("pubkey-metadata")
	for _, r := range oxPubKeyListXMLPubsubItemsItemPklPm {
		date := r.SelectAttr("date")
		if date == nil {
			continue
		}
		fingerprint := r.SelectAttr("v4-fingerprint")
		if fingerprint == nil {
			continue
		}
		keyDate, err := time.Parse(time.RFC3339, date.Value)
		if err != nil {
			return nil, err
		}
		if keyDate.After(newestKey) {
			newestKey = keyDate
			pubKeyRingID = fingerprint.Value
		}
	}
	if pubKeyRingID == "none" {
		return nil, errors.New("server didn't provide public key fingerprints for " + recipient)
	}

	pubKeyRingLocation, err := oxGetPubKeyLoc(pubKeyRingID)
	if err != nil {
		return nil, err
	}

	pubKeyReadXML := etree.NewDocument()
	err = pubKeyReadXML.ReadFromFile(pubKeyRingLocation)
	if err == nil {
		date := pubKeyReadXML.SelectElement("date")
		if date != nil {
			savedKeysDate, err := time.Parse(time.RFC3339, date.Text())
			if err != nil {
				return nil, err
			}
			if !savedKeysDate.Before(newestKey) {
				pubKeys := pubKeyReadXML.SelectElements("pubkey")
				if pubKeys == nil {
					return nil, errors.New("couldn't read public keys from cache")
				}
				for _, r := range pubKeys {
					keyByte, err := base64.StdEncoding.DecodeString(r.Text())
					if err != nil {
						return nil, err
					}
					key, err := crypto.NewKey(keyByte)
					if err != nil {
						return nil, err
					}
					if !key.IsExpired() {
						err = publicKeyRing.AddKey(key)
						if err != nil {
							return nil, err
						}
					}
				}
				if publicKeyRing.CanEncrypt() {
					return publicKeyRing, nil
				}
			}
		}
	}
	pubKeyRing, err := oxRecvPublicKeys(client, iqc, recipient, pubKeyRingID)
	if err != nil {
		return nil, err
	}
	pubKeySaveXML := etree.NewDocument()
	date := pubKeySaveXML.CreateElement("date")
	date.SetText(newestKey.Format(time.RFC3339))
	for _, key := range pubKeyRing.GetKeys() {
		keySerialized, err := key.Serialize()
		if err != nil {
			return nil, err
		}
		if err != nil {
			return nil, err
		}
		saveKey := pubKeySaveXML.CreateElement("pubkey")
		saveKey.SetText(base64.StdEncoding.EncodeToString(keySerialized))
	}
	err = pubKeySaveXML.WriteToFile(pubKeyRingLocation)
	if err != nil {
		return nil, err
	}
	return pubKeyRing, nil
}

func oxEncrypt(client *xmpp.Client, iqc chan xmpp.IQ, oxPrivKey *crypto.Key,
	recipient string, keyRing *crypto.KeyRing, message string) (string, error) {
	if message == "" {
		return "", nil
	}
	privKeyRing, err := crypto.NewKeyRing(oxPrivKey)
	if err != nil {
		return "error", err
	}
	ownJid := strings.Split(client.JID(), "/")[0]
	if recipient != ownJid {
		opk, err := oxPrivKey.GetPublicKey()
		if err == nil {
			ownKey, _ := crypto.NewKey(opk)
			_ = keyRing.AddKey(ownKey)
		}
	}
	ocm, err := oxCryptMessage(recipient, message)
	if err != nil {
		return "error", err
	}
	plainMessage := crypto.NewPlainMessage([]byte(ocm))
	pgpMessage, err := keyRing.Encrypt(plainMessage, privKeyRing)
	if err != nil {
		return "error", err
	}
	oms, err := oxMessage(recipient, pgpMessage)
	if err != nil {
		return "error", err
	}
	return oms, nil
}