package omemo

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/xml"
	"math/big"
)

// Minimal XEP-0420 Stanza Content Encryption envelope, used only as the
// plaintext OMEMO 2 (modern) content encryption operates on - just enough
// to carry a plain message body, per the spec's requirement that the
// envelope MUST contain <rpad/> and SHOULD contain <from/> (MUST contain
// <to/> for MUC, which this gateway doesn't do OMEMO for - see the plan's
// scope note on MUC).
type sceEnvelope struct {
	XMLName xml.Name    `xml:"urn:xmpp:sce:1 envelope"`
	Content sceContent  `xml:"content"`
	RPad    string      `xml:"rpad"`
	From    *sceJIDAttr `xml:"from"`
}

type sceContent struct {
	Body sceBody `xml:"body"`
}

type sceBody struct {
	XMLName xml.Name `xml:"jabber:client body"`
	Text    string   `xml:",chardata"`
}

type sceJIDAttr struct {
	JID string `xml:"jid,attr"`
}

// sceEncode wraps plaintext in an SCE envelope and serializes it - the
// "plaintext" that modern OMEMO's outer content encryption actually
// operates on. fromJID becomes the <from/> binding (empty to omit it).
func sceEncode(plaintext []byte, fromJID string) ([]byte, error) {
	rpad, err := randomPadding()
	if err != nil {
		return nil, err
	}
	env := sceEnvelope{
		Content: sceContent{Body: sceBody{Text: string(plaintext)}},
		RPad:    rpad,
	}
	if fromJID != "" {
		env.From = &sceJIDAttr{JID: fromJID}
	}
	return xml.Marshal(env)
}

// sceDecode reverses sceEncode, extracting the plain body text.
func sceDecode(data []byte) ([]byte, error) {
	var env sceEnvelope
	if err := xml.Unmarshal(data, &env); err != nil {
		return nil, err
	}
	return []byte(env.Content.Body.Text), nil
}

// randomPadding returns base64-encoded random padding of a random length
// (1-63 bytes before encoding) to mask plaintext length, per XEP-0420's
// MUST-contain-<rpad/> requirement.
func randomPadding() (string, error) {
	n, err := rand.Int(rand.Reader, big.NewInt(63))
	if err != nil {
		return "", err
	}
	buf := make([]byte, n.Int64()+1)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}