fckeuspy-go/vault_service.go

package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	encrypt "fckeuspy-go/lib"
	"fmt"
	mrand "math/rand"
	"sync"
	"time"
)

// VaultService implementuje ServiceFacade nad SecureJSONStore.
// Využívá existující hybridní šifrování z encrypt.Service (re-use kódu voláním helperů?)
// Pro jednoduchost zopakujeme potřebnou část: Encrypt/Decrypt využijí dočasný encrypt.Service
// inicializovaný z načtených klíčů.
type VaultService struct {
	store   encrypt.SecureJSONStore
	priv    *rsa.PrivateKey
	pubPEM  string
	certPEM string
	chatMu  sync.Mutex
}

func NewVaultService(store encrypt.SecureJSONStore) (*VaultService, error) {
	vs := &VaultService{store: store}
	if err := vs.loadIdentity(); err != nil {
		return nil, err
	}
	return vs, nil
}

func (v *VaultService) loadIdentity() error {
	privPem := v.store.IdentityPrivatePEM()
	if privPem == "" {
		return errors.New("missing private key in store")
	}
	block, _ := pem.Decode([]byte(privPem))
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return errors.New("invalid private key PEM")
	}
	key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return err
	}
	v.priv = key
	v.pubPEM = v.store.IdentityPublicPEM()
	v.certPEM = v.store.IdentityCertPEM()
	if v.pubPEM == "" { // derive public if missing
		pubASN1, _ := x509.MarshalPKIXPublicKey(&key.PublicKey)
		v.pubPEM = string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubASN1}))
	}
	return nil
}

func (v *VaultService) PublicPEM() string  { return v.pubPEM }
func (v *VaultService) PublicCert() string { return v.certPEM }

// Encrypt provede hybridní šifrování stejně jako původní Service.Encrypt.
func (v *VaultService) Encrypt(message, peerPEMorCert string) (string, error) {
	// Vytvoříme ad-hoc Service s privátním klíčem jen pro volání existující logiky? Nejjednodušší je zkopírovat potřebnou část kódu.
	return encryptHybrid(v.priv, message, peerPEMorCert)
}

func (v *VaultService) EncryptChat(message, peerPEMorCert string) (string, error) {
	return encryptChatHybrid(v.priv, message, peerPEMorCert)
}

// Decrypt provede rozšifrování.
func (v *VaultService) Decrypt(payload string) (string, error) { return decryptHybrid(v.priv, payload) }

func (v *VaultService) DecryptChat(payload string) (DecryptedChatMessage, error) {
	return decryptChatHybrid(v.priv, payload)
}

// --- Lokální helpery (duplikace z encrypt.Service, zredukované) ---

type hybridEnvelope struct {
	EK     string `json:"ek"`
	N      string `json:"n"`
	CT     string `json:"ct"`
	SelfEK string `json:"sek,omitempty"`
	SentAt string `json:"at,omitempty"`
}

// --- Contacts management ---

type Contact struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	Cert string `json:"cert"` // full PEM cert (or public key PEM)
}

const contactsKey = "contacts"

const (
	chatHistoryKey            = "chat_history"
	persistedChatDirectionIn  = "incoming"
	persistedChatDirectionOut = "outgoing"
)

type ChatMessage struct {
	Direction   string    `json:"direction"`
	Text        string    `json:"text"`
	CreatedAt   time.Time `json:"createdAt"`
	PayloadHash string    `json:"payloadHash"`
}

type DecryptedChatMessage struct {
	Direction string
	Text      string
	CreatedAt time.Time
}

func (v *VaultService) ListContacts() ([]Contact, error) {
	var list []Contact
	if !v.store.Has(contactsKey) {
		return []Contact{}, nil
	}
	if err := v.store.Get(contactsKey, &list); err != nil {
		return nil, err
	}
	return list, nil
}

func (v *VaultService) ListChatMessages() ([]ChatMessage, error) {
	v.chatMu.Lock()
	defer v.chatMu.Unlock()

	return v.listChatMessagesLocked()
}

func (v *VaultService) listChatMessagesLocked() ([]ChatMessage, error) {
	if !v.store.Has(chatHistoryKey) {
		return []ChatMessage{}, nil
	}

	var messages []ChatMessage
	if err := v.store.Get(chatHistoryKey, &messages); err != nil {
		return nil, err
	}
	return messages, nil
}

func (v *VaultService) AppendChatMessage(direction, text, payload string) (ChatMessage, bool, error) {
	return v.appendChatMessage(direction, text, payload, time.Now().UTC())
}

func (v *VaultService) AppendScannedChatMessage(direction, text, payload string, createdAt time.Time) (ChatMessage, bool, error) {
	if createdAt.IsZero() {
		createdAt = time.Now().UTC()
	}
	return v.appendChatMessage(direction, text, payload, createdAt)
}

func (v *VaultService) appendChatMessage(direction, text, payload string, createdAt time.Time) (ChatMessage, bool, error) {
	if direction != persistedChatDirectionIn && direction != persistedChatDirectionOut {
		return ChatMessage{}, false, fmt.Errorf("invalid chat message direction %q", direction)
	}
	if payload == "" {
		return ChatMessage{}, false, errors.New("empty chat payload")
	}

	payloadSum := sha256.Sum256([]byte(payload))
	payloadHash := fmt.Sprintf("%x", payloadSum[:])

	v.chatMu.Lock()
	defer v.chatMu.Unlock()

	messages, err := v.listChatMessagesLocked()
	if err != nil {
		return ChatMessage{}, false, err
	}
	for _, message := range messages {
		if message.PayloadHash == payloadHash {
			return message, false, nil
		}
	}

	message := ChatMessage{
		Direction:   direction,
		Text:        text,
		CreatedAt:   createdAt.UTC(),
		PayloadHash: payloadHash,
	}
	messages = append(messages, message)
	if err := v.store.Put(chatHistoryKey, messages); err != nil {
		return ChatMessage{}, false, err
	}
	if err := v.store.Flush(); err != nil {
		return ChatMessage{}, false, err
	}
	return message, true, nil
}

func (v *VaultService) SaveContact(c Contact) error {
	list, _ := v.ListContacts()
	// upsert by ID; if empty ID, assign a new unique random ID to avoid overwriting
	if c.ID == "" {
		c.ID = generateUniqueContactID(list)
	}
	replaced := false
	for i := range list {
		if list[i].ID == c.ID {
			list[i] = c
			replaced = true
			break
		}
	}
	if !replaced {
		list = append(list, c)
	}
	if err := v.store.Put(contactsKey, list); err != nil {
		return err
	}
	return v.store.Flush()
}

func (v *VaultService) DeleteContact(id string) error {
	list, _ := v.ListContacts()
	out := list[:0]
	for _, c := range list {
		if c.ID != id {
			out = append(out, c)
		}
	}
	if err := v.store.Put(contactsKey, out); err != nil {
		return err
	}
	return v.store.Flush()
}

func deriveContactID(c Contact) string {
	// simple stable ID: prefer CN from cert; fallback to sha256 of cert
	if cn := extractCN(c.Cert); cn != "" {
		return cn
	}
	sum := sha256.Sum256([]byte(c.Cert))
	return fmt.Sprintf("%x", sum[:8])
}

// generateUniqueContactID creates a short random hex ID that does not collide with existing list.
func generateUniqueContactID(existing []Contact) string {
	exists := make(map[string]struct{}, len(existing))
	for _, c := range existing {
		exists[c.ID] = struct{}{}
	}
	for {
		b := make([]byte, 6)
		if _, err := rand.Read(b); err != nil {
			// fallback to time-based seed rand
			n := mrand.Int63()
			return fmt.Sprintf("%x", n)
		}
		id := fmt.Sprintf("%x", b)
		if _, ok := exists[id]; !ok {
			return id
		}
	}
}

func extractCN(pemText string) string {
	// Support multiple concatenated PEM blocks (public key + cert, or chain)
	rest := []byte(pemText)
	for len(rest) > 0 {
		var block *pem.Block
		block, rest = pem.Decode(rest)
		if block == nil {
			break
		}
		if block.Type == "CERTIFICATE" {
			if cert, err := x509.ParseCertificate(block.Bytes); err == nil {
				return cert.Subject.CommonName
			}
		}
		// continue to next block until certificate found
	}
	return ""
}

func encryptHybrid(priv *rsa.PrivateKey, message, peerPEMorCert string) (string, error) {
	return encryptHybridEnvelope(nil, message, peerPEMorCert)
}

func encryptChatHybrid(priv *rsa.PrivateKey, message, peerPEMorCert string) (string, error) {
	if priv == nil {
		return "", errors.New("missing private key")
	}
	return encryptHybridEnvelope(&priv.PublicKey, message, peerPEMorCert)
}

func encryptHybridEnvelope(selfPublicKey *rsa.PublicKey, message, peerPEMorCert string) (string, error) {
	pubKey, err := encrypt.ParsePeerPublicKey(peerPEMorCert)
	if err != nil {
		return "", err
	}
	aesKey := make([]byte, 32)
	if _, err := rand.Read(aesKey); err != nil {
		return "", err
	}
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	ct := gcm.Seal(nil, nonce, []byte(message), nil)
	ek, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubKey, aesKey, []byte{})
	if err != nil {
		return "", err
	}
	env := hybridEnvelope{
		EK: base64.StdEncoding.EncodeToString(ek),
		N:  base64.StdEncoding.EncodeToString(nonce),
		CT: base64.StdEncoding.EncodeToString(ct),
	}
	if selfPublicKey != nil {
		selfEK, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, selfPublicKey, aesKey, []byte{})
		if err != nil {
			return "", err
		}
		env.SelfEK = base64.StdEncoding.EncodeToString(selfEK)
		env.SentAt = time.Now().UTC().Format(time.RFC3339Nano)
	}
	out, _ := json.MarshalIndent(env, "", "  ")
	return string(out), nil
}

func decryptHybrid(priv *rsa.PrivateKey, payload string) (string, error) {
	env, nonce, ct, err := decodeHybridEnvelope(payload)
	if err != nil {
		return "", err
	}
	return decryptHybridPayload(priv, env.EK, nonce, ct)
}

func decryptChatHybrid(priv *rsa.PrivateKey, payload string) (DecryptedChatMessage, error) {
	env, nonce, ct, err := decodeHybridEnvelope(payload)
	if err != nil {
		return DecryptedChatMessage{}, err
	}
	createdAt := time.Time{}
	if env.SentAt != "" {
		createdAt, _ = time.Parse(time.RFC3339Nano, env.SentAt)
	}
	if env.SelfEK != "" {
		plaintext, selfErr := decryptHybridPayload(priv, env.SelfEK, nonce, ct)
		if selfErr == nil {
			return DecryptedChatMessage{Direction: persistedChatDirectionOut, Text: plaintext, CreatedAt: createdAt}, nil
		}
	}
	plaintext, err := decryptHybridPayload(priv, env.EK, nonce, ct)
	if err != nil {
		return DecryptedChatMessage{}, err
	}
	return DecryptedChatMessage{Direction: persistedChatDirectionIn, Text: plaintext, CreatedAt: createdAt}, nil
}

func decodeHybridEnvelope(payload string) (hybridEnvelope, []byte, []byte, error) {
	var env hybridEnvelope
	if err := json.Unmarshal([]byte(payload), &env); err != nil {
		return hybridEnvelope{}, nil, nil, fmt.Errorf("invalid JSON: %w", err)
	}
	nonce, err := base64.StdEncoding.DecodeString(env.N)
	if err != nil {
		return hybridEnvelope{}, nil, nil, fmt.Errorf("n b64: %w", err)
	}
	ct, err := base64.StdEncoding.DecodeString(env.CT)
	if err != nil {
		return hybridEnvelope{}, nil, nil, fmt.Errorf("ct b64: %w", err)
	}
	return env, nonce, ct, nil
}

func decryptHybridPayload(priv *rsa.PrivateKey, encodedKey string, nonce, ct []byte) (string, error) {
	if priv == nil {
		return "", errors.New("missing private key")
	}
	ek, err := base64.StdEncoding.DecodeString(encodedKey)
	if err != nil {
		return "", fmt.Errorf("ek b64: %w", err)
	}
	aesKey, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, priv, ek, []byte{})
	if err != nil {
		return "", fmt.Errorf("RSA-OAEP decrypt: %w", err)
	}
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	pt, err := gcm.Open(nil, nonce, ct, nil)
	if err != nil {
		return "", fmt.Errorf("GCM open: %w", err)
	}
	return string(pt), nil
}