// 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"
	"encoding/xml"
	"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 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 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("Couldn't unlock private key.")
		}
	}
	isLocked, err := key.IsLocked()
	if err != nil {
		return nil, err
	}
	if isLocked {
		log.Fatal("Private key is locked.")
	}
	if key.IsExpired() {
		return nil, errors.New("Key is expired: " + key.GetFingerprint())
	}
	return key, nil
}

func oxStorePrivKey(jid string, privKey string) error {
	dataFile, err := oxGetPrivKeyLoc(jid)
	if err != nil {
		log.Fatal(err)
	}
	var file *os.File
	if _, err := os.Stat(dataFile); os.IsNotExist(err) {
		file, err = os.Create(dataFile)
		if err != nil {
			return err
		}
	} else {
		file, err = os.OpenFile(dataFile, os.O_RDWR, 0600)
		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(privKey))
	if err != nil {
		return err
	}

	return nil
}

func oxGenPrivKey(jid string, client *xmpp.Client, passphrase string) error {
	xmppUri := "xmpp:" + jid
	key, err := crypto.GenerateKey(xmppUri, xmppUri, "x25519", 0)
	if err != nil {
		return err
	}
	decodedPubKey, err := key.GetPublicKey()
	if err != nil {
		return err
	}
	pubKeyBase64 := base64.StdEncoding.EncodeToString(decodedPubKey)
	if passphrase != "" {
		key, err = key.Lock([]byte(passphrase))
		if err != nil {
			return err
		}
	}
	keySerialized, err := key.Serialize()
	if err != nil {
		return err
	}
	err = oxStorePrivKey(jid,
		base64.StdEncoding.EncodeToString(keySerialized))
	if err != nil {
		log.Fatal(err)
	}
	keyCreated := time.Now().UTC().Format("2006-01-02T15:04:05Z")
	pubKey, err := crypto.NewKey(decodedPubKey)
	if err != nil {
		return err
	}
	fingerprint := strings.ToUpper(pubKey.GetFingerprint())
	root := etree.NewDocument()
	pubsub := root.CreateElement("pubsub")
	pubsub.CreateAttr("xmlns", nsPubsub)
	publish := pubsub.CreateElement("publish")
	publish.CreateAttr("node", nsOxPubKeys+":"+fingerprint)
	item := publish.CreateElement("item")
	item.CreateAttr("id", keyCreated)
	pubkey := item.CreateElement("pubkey")
	pubkey.CreateAttr("xmlns", nsOx)
	data := pubkey.CreateElement("data")
	data.CreateText(pubKeyBase64)
	publishoptions := pubsub.CreateElement("publish-options")
	x := publishoptions.CreateElement("x")
	x.CreateAttr("xmlns", nsJabberData)
	x.CreateAttr("type", "submit")
	field := x.CreateElement("field")
	field.CreateAttr("var", "FORM_TYPE")
	field.CreateAttr("type", "hidden")
	value := field.CreateElement("value")
	value.CreateText(pubsubPubOptions)
	field = x.CreateElement("field")
	field.CreateAttr("var", "pubsub#access_model")
	value = field.CreateElement("value")
	value.CreateText("open")
	xmlstring, err := root.WriteToString()
	if err != nil {
		return err
	}
	iqReply, err := sendIQ(client, jid, "set", xmlstring)
	if err != nil {
		return err
	}
	if iqReply.Type != "result" {
		return errors.New("Error while publishing public key.")
	}

	ownPubKeyFromPubsub, err := oxRecvPublicKey(client, jid, fingerprint)
	if err != nil {
		return errors.New("Couldn't successfully verify public key upload.")
	}
	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.")
	}

	root = etree.NewDocument()
	pubsub = root.CreateElement("pubsub")
	pubsub.CreateAttr("xmlns", nsPubsub)
	publish = pubsub.CreateElement("publish")
	publish.CreateAttr("node", nsOxPubKeys)
	item = publish.CreateElement("item")
	pubkeyslist := item.CreateElement("public-keys-list")
	pubkeyslist.CreateAttr("xmlns", nsOx)
	pubkeymeta := pubkeyslist.CreateElement("pubkey-metadata")
	pubkeymeta.CreateAttr("v4-fingerprint", fingerprint)
	pubkeymeta.CreateAttr("date", keyCreated)
	publishoptions = pubsub.CreateElement("publish-options")
	x = publishoptions.CreateElement("x")
	x.CreateAttr("xmlns", nsJabberData)
	x.CreateAttr("type", "submit")
	field = x.CreateElement("field")
	field.CreateAttr("var", "FORM_TYPE")
	field.CreateAttr("type", "hidden")
	value = field.CreateElement("value")
	value.CreateText(pubsubPubOptions)
	field = x.CreateElement("field")
	field.CreateAttr("var", "pubsub#access_model")
	value = field.CreateElement("value")
	value.CreateText("open")
	xmlstring, err = root.WriteToString()
	if err != nil {
		return err
	}
	iqReply, err = sendIQ(client, jid, "set", xmlstring)
	if err != nil {
		return err
	}
	if iqReply.Type != "result" {
		return errors.New("Couldn't publish public key list.")
	}
	return nil
}

func oxRecvPublicKey(client *xmpp.Client, recipient string,
	fingerprint string) (*crypto.Key, error) {
	var oxPublicKeyRequest IQPubsubRequest
	var oxPublicKeyXML OxPublicKey
	oxPublicKeyRequest.Xmlns = nsPubsub
	oxPublicKeyRequest.Items.Node = nsOxPubKeys + ":" + fingerprint
	oxPublicKeyRequest.Items.MaxItems = "1"
	opk, err := xml.Marshal(oxPublicKeyRequest)
	if err != nil {
		return nil, err
	}
	oxPublicKey, err := sendIQ(client, recipient, "get", string(opk))
	if err != nil {
		return nil, err
	}
	if oxPublicKey.Type != "result" {
		return nil, errors.New("Error while requesting public key for " +
			recipient)
	}
	err = xml.Unmarshal(oxPublicKey.Query, &oxPublicKeyXML)
	if err != nil {
		return nil, err
	}

	decodedPubKey, err := base64.StdEncoding.DecodeString(oxPublicKeyXML.Items.Item.Pubkey.Data)
	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)
	}
	return key, nil
}

func oxGetPublicKey(client *xmpp.Client, recipient string) (*crypto.Key, error) {
	var oxPublicKeyListRequest IQPubsubRequest
	var oxPublicKeyListXML OxPublicKeysList
	oxPublicKeyListRequest.Xmlns = nsPubsub
	oxPublicKeyListRequest.Items.Node = nsOxPubKeys
	oxPublicKeyListRequest.Items.MaxItems = "100"
	opkl, err := xml.Marshal(oxPublicKeyListRequest)
	if err != nil {
		log.Fatal(err)
	}
	oxPublicKeyList, err := sendIQ(client, recipient, "get", string(opkl))
	if err != nil {
		log.Fatal(err)
	}
	if oxPublicKeyList.Type != "result" {
		return nil, errors.New("Error while requesting public openpgp keys for " +
			recipient)
	}
	err = xml.Unmarshal(oxPublicKeyList.Query, &oxPublicKeyListXML)
	if err != nil {
		return nil, err
	}

	fingerprint := "none"
	newestKey, err := time.Parse(time.RFC3339, "2006-01-02T15:04:05Z")
	if err != nil {
		return nil, err
	}
	for _, r := range oxPublicKeyListXML.Items.Item.PublicKeysList.PubkeyMetadata {
		keyDate, err := time.Parse(time.RFC3339, r.Date)
		if err != nil {
			return nil, err
		}
		if keyDate.After(newestKey) {
			newestKey = keyDate
			fingerprint = r.V4Fingerprint
		}
	}
	if fingerprint == "none" {
		return nil, errors.New("server didn't provide public key fingerprints for " + recipient)
	}

	key, err := oxRecvPublicKey(client, recipient, fingerprint)
	if err != nil {
		return nil, errors.New("Couldn't fetch public key.")
	}

	return key, nil
}

func oxEncrypt(client *xmpp.Client, oxPrivKey *crypto.Key, recipient string,
	recipientKey *crypto.Key, message string) (string, error) {
	var oxCryptMessage OxCryptElement
	var oxMessage OxMessageElement
	keyRing, err := crypto.NewKeyRing(recipientKey)
	if err != nil {
		return "error", err
	}
	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)
		}
	}
	oxCryptMessage.Xmlns = nsOx
	oxCryptMessage.To.Jid = recipient
	oxCryptMessage.Time.Stamp = time.Now().UTC().Format("2006-01-02T15:04:05Z")
	oxCryptMessage.Rpad = getRpad()
	oxCryptMessage.Payload.Body.Xmlns = nsJabberClient
	oxCryptMessage.Payload.Body.Text = message
	ocm, err := xml.Marshal(oxCryptMessage)
	if err != nil {
		return "error", err
	}
	plainMessage := crypto.NewPlainMessage([]byte(ocm))
	pgpMessage, err := keyRing.Encrypt(plainMessage, privKeyRing)
	if err != nil {
		return "error", err
	}
	oxMessage.To = recipient
	oxMessage.Id = getID()
	oxMessage.From = client.JID()
	oxMessage.Openpgp.Text = base64.StdEncoding.EncodeToString(pgpMessage.Data)
	oxMessage.Openpgp.Xmlns = nsOx
	oxMessage.Encryption.Xmlns = nsEme
	oxMessage.Encryption.Namespace = nsOx
	oxMessage.Body = oxAltBody
	om, err := xml.Marshal(oxMessage)
	if err != nil {
		return "error", err
	}

	return string(om), nil
}