package ankersolix

// Payload encryption — the mobile app's "algo_ecdh" mode for the REST API.
//
// Anker's servers currently accept plain JSON bodies, and that is what this
// connector has always sent. The app itself does not: it runs a key exchange
// once per login and then encrypts every request body and reads back every
// response body encrypted, with an HMAC over each request. anker-solix-api
// v3.8.2 worked the scheme out end to end (session.py, AnkerEncryptionHandler),
// and this file is that handler in Go, kept behind the payloadEncryption config
// switch — off by default, as upstream leaves it, because plain bodies are what
// the cloud has been answering for years and any slip in a port of the cipher
// would take every read down at once. It exists so the day Anker stops
// tolerating plain bodies is a config change, not an outage.
//
// The scheme, in the order it runs:
//
//  1. Handshake. A fresh P-256 key pair is generated; its uncompressed point
//     (0x04 || X || Y) is taken as its 130-character ascii hex and wrapped —
//     AES-128-CBC under the region's presetKey with a random IV prepended, then
//     base64 — and POSTed as client_public_key to openapi/oauth/key/exchange with
//     the usual auth headers. The x-signature on that request is
//     HMAC-SHA256(key = the presetKey's ascii hex, msg = ts+once+client_public_key).
//     The reply's server_public_key is wrapped the same way; unwrapping it and
//     running ECDH gives the 32-byte shared secret. There is no KDF.
//  2. Requests. The body is compact JSON, AES-128-CBC under secret[:16] (the
//     "securityKey"), random IV prepended, PKCS#7, base64, sent as text/plain.
//     x-signature is HMAC-SHA256(key = ascii hex of secret[:16],
//     msg = ts+once+body). x-key-ident is fixed for the life of the key pair
//     and is how the server finds the secret again; x-request-once is fresh
//     per request.
//  3. Responses come back with data as a base64 string and a signature field
//     beside it; data decrypts with the same key and cipher to the JSON the
//     plain mode would have received.
//
// The presetKey is an app constant, identical for both regional servers today
// (upstream's API_PRESET_KEYS), not a user secret.

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
)

const (
	epKeyExchange = "openapi/oauth/key/exchange"

	// presetKeyHex is the region AES-128 key the app carries in its stored
	// userConfig as "presetKey". It wraps the two public keys of the handshake and
	// keys the handshake's signature. EU and global use the same one.
	presetKeyHex = "e8ad18f61bbd3fbd52d5ed12d14d3b9c"
)

// payloadCipher is one token's encryption state: the key pair the handshake was
// run with, the identifier the server files its result under, and the secret
// the handshake produced. It lives beside the token in the session and is
// discarded with it — a new login is a new handshake.
type payloadCipher struct {
	priv     *ecdh.PrivateKey
	keyIdent string // x-key-ident: 32 hex chars, random, constant per key pair
	preset   []byte // the region presetKey, raw
	secret   []byte // ECDH shared secret; nil until the exchange has answered
}

// newPayloadCipher generates the key pair and identifier for one handshake.
func newPayloadCipher() (*payloadCipher, error) {
	priv, err := ecdh.P256().GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("anker-solix: generate key: %w", err)
	}
	preset, err := hex.DecodeString(presetKeyHex)
	if err != nil {
		return nil, fmt.Errorf("anker-solix: preset key: %w", err)
	}
	ident := make([]byte, 16)
	if _, err := rand.Read(ident); err != nil {
		return nil, fmt.Errorf("anker-solix: key ident: %w", err)
	}
	return &payloadCipher{priv: priv, keyIdent: hex.EncodeToString(ident), preset: preset}, nil
}

// ready reports whether the handshake has produced a secret.
func (c *payloadCipher) ready() bool { return c != nil && c.secret != nil }

// clientPublicKey is the handshake body's one field: the uncompressed point as
// ascii hex, wrapped under the presetKey.
func (c *payloadCipher) clientPublicKey() (string, error) {
	return aesCBCSeal(c.preset, []byte(hex.EncodeToString(c.priv.PublicKey().Bytes())))
}

// acceptServerKey unwraps the handshake reply's server_public_key and derives
// the shared secret from it.
func (c *payloadCipher) acceptServerKey(wrapped string) error {
	pointHex, err := aesCBCOpen(c.preset, wrapped)
	if err != nil {
		return fmt.Errorf("unwrap server key: %w", err)
	}
	point, err := hex.DecodeString(strings.TrimSpace(string(pointHex)))
	if err != nil {
		return fmt.Errorf("server key is not hex: %w", err)
	}
	pub, err := ecdh.P256().NewPublicKey(point)
	if err != nil {
		return fmt.Errorf("server key is not a P-256 point: %w", err)
	}
	secret, err := c.priv.ECDH(pub)
	if err != nil {
		return fmt.Errorf("derive shared secret: %w", err)
	}
	c.secret = secret
	return nil
}

// bodyKey is the AES-128 key both directions use after the handshake.
func (c *payloadCipher) bodyKey() []byte { return c.secret[:16] }

// sealBody encrypts a request payload for the wire. GETs, which carry no
// payload, seal an empty object — what the reference implementation sends.
func (c *payloadCipher) sealBody(payload map[string]any) (string, error) {
	if payload == nil {
		payload = map[string]any{}
	}
	plain, err := json.Marshal(payload)
	if err != nil {
		return "", fmt.Errorf("anker-solix: encode request: %w", err)
	}
	return aesCBCSeal(c.bodyKey(), plain)
}

// signRequest sets the encryption headers on a request whose body is already
// on the wire form given. The handshake signs under the presetKey and keeps a
// JSON content type; every request after it signs under the session key and
// goes out as text/plain.
func (c *payloadCipher) signRequest(req *http.Request, ts, wire string, handshake bool) {
	var signKey string
	if handshake {
		signKey = presetKeyHex
		req.Header.Set("content-type", "application/json")
		req.Header.Set("x-auth-ts", ts)
	} else {
		signKey = hex.EncodeToString(c.bodyKey())
		req.Header.Set("content-type", "text/plain")
	}
	once := make([]byte, 16)
	_, _ = rand.Read(once)
	onceHex := hex.EncodeToString(once)

	mac := hmac.New(sha256.New, []byte(signKey))
	mac.Write([]byte(ts + "+" + onceHex + "+" + wire))

	req.Header.Set("x-app-key", "")
	req.Header.Set("x-encryption-info", "algo_ecdh")
	req.Header.Set("x-replay-info", "replay")
	req.Header.Set("x-key-ident", c.keyIdent)
	req.Header.Set("x-request-once", onceHex)
	req.Header.Set("x-request-ts", ts)
	req.Header.Set("x-signature", hex.EncodeToString(mac.Sum(nil)))
}

// unsealEnvelope turns an encrypted response into the envelope the plain mode
// receives: when a signature field is present and data is a string, data is
// decrypted and put back as the JSON it holds. Anything else — an error
// envelope, a plain reply — is returned as it came.
func (c *payloadCipher) unsealEnvelope(body []byte) ([]byte, error) {
	var env map[string]json.RawMessage
	if err := json.Unmarshal(body, &env); err != nil {
		return body, nil
	}
	if _, signed := env["signature"]; !signed {
		return body, nil
	}
	var wrapped string
	if err := json.Unmarshal(env["data"], &wrapped); err != nil || wrapped == "" {
		return body, nil
	}
	plain, err := aesCBCOpen(c.bodyKey(), wrapped)
	if err != nil {
		return nil, fmt.Errorf("anker-solix: decrypt response: %w", err)
	}
	if !json.Valid(plain) {
		return nil, errors.New("anker-solix: decrypted response is not JSON")
	}
	env["data"] = json.RawMessage(plain)
	return json.Marshal(env)
}

// aesCBCSeal is the one envelope every wrapped value in the scheme uses:
// AES-CBC under key, a random IV prepended to the ciphertext, PKCS#7 padding,
// standard base64. The key's length picks AES-128 or -256.
func aesCBCSeal(key, plain []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("aes cipher: %w", err)
	}
	iv := make([]byte, block.BlockSize())
	if _, err := rand.Read(iv); err != nil {
		return "", fmt.Errorf("iv: %w", err)
	}
	padded := pkcs7Pad(plain, block.BlockSize())
	out := make([]byte, len(iv)+len(padded))
	copy(out, iv)
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(out[len(iv):], padded)
	return base64.StdEncoding.EncodeToString(out), nil
}

// aesCBCOpen reverses aesCBCSeal.
func aesCBCOpen(key []byte, wrapped string) ([]byte, error) {
	raw, err := base64.StdEncoding.DecodeString(strings.TrimSpace(wrapped))
	if err != nil {
		return nil, fmt.Errorf("base64: %w", err)
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("aes cipher: %w", err)
	}
	bs := block.BlockSize()
	if len(raw) < 2*bs || len(raw)%bs != 0 {
		return nil, fmt.Errorf("ciphertext length %d is not iv plus whole blocks", len(raw))
	}
	iv, ct := raw[:bs], raw[bs:]
	plain := make([]byte, len(ct))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(plain, ct)
	return pkcs7Unpad(plain, bs)
}

// pkcs7Unpad strips the padding pkcs7Pad added, refusing anything malformed.
func pkcs7Unpad(data []byte, blockSize int) ([]byte, error) {
	if len(data) == 0 || len(data)%blockSize != 0 {
		return nil, errors.New("bad padding: length")
	}
	pad := int(data[len(data)-1])
	if pad == 0 || pad > blockSize || pad > len(data) {
		return nil, errors.New("bad padding: count")
	}
	if !bytes.Equal(data[len(data)-pad:], bytes.Repeat([]byte{byte(pad)}, pad)) {
		return nil, errors.New("bad padding: bytes")
	}
	return data[:len(data)-pad], nil
}