Necropolis-C2/pkg/transport/messenger_test.go

package transport

import (
	"bytes"
	"context"
	"strings"
	"testing"
	"time"

	"github.com/libp2p/go-libp2p"

	"github.com/Yenn503/NecropolisC2/pkg/cryptography"
)

// minimalHost creates a bare libp2p host for unit tests. No DHT, no relay.
func minimalHost(t *testing.T) (*Node, context.Context, context.CancelFunc) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	h, err := libp2p.New(libp2p.ListenAddrStrings("/ip4/127.0.0.1/tcp/0"))
	if err != nil {
		t.Fatalf("create host: %v", err)
	}
	t.Cleanup(func() { h.Close(); cancel() })
	return &Node{Host: h, ctx: ctx, cancel: func() {}}, ctx, cancel
}

func TestCreateEnvelopeStoresToken(t *testing.T) {
	op, err := cryptography.GenerateOperatorKey()
	if err != nil {
		t.Fatalf("generate operator key: %v", err)
	}
	n, ctx, _ := minimalHost(t)
	_ = ctx
	m := NewOperatorMessenger(ctx, n, op)
	m.SetAuthToken(op.AuthToken)
	env := m.CreateEnvelope(MsgTypePs, []byte("test"))

	if len(env.Token) != 32 {
		t.Fatalf("token missing from envelope: got %d bytes, want 32", len(env.Token))
	}
	if !bytes.Equal(env.Token, op.AuthToken) {
		t.Fatal("envelope token != operator auth token")
	}
	if env.ID == 0 {
		t.Fatal("envelope ID is zero — random generation failed")
	}
}

func TestReplayProtection(t *testing.T) {
	op, err := cryptography.GenerateOperatorKey()
	if err != nil {
		t.Fatalf("generate operator key: %v", err)
	}
	n, ctx, _ := minimalHost(t)
	m := NewOperatorMessenger(ctx, n, op)

	env1 := m.CreateEnvelope(MsgTypePs, []byte("test"))
	env2 := m.CreateEnvelope(MsgTypePs, []byte("test"))
	env2.ID = env1.ID // force same ID

	if m.IsReplay(env1.ID) {
		t.Fatal("first envelope flagged as replay")
	}
	if !m.IsReplay(env1.ID) {
		t.Fatal("same ID not detected as replay on second check")
	}
	if !m.IsReplay(env2.ID) {
		t.Fatal("replay envelope with same ID not detected")
	}
}

func TestAuthTokenRejection(t *testing.T) {
	op, err := cryptography.GenerateOperatorKey()
	if err != nil {
		t.Fatalf("generate operator key: %v", err)
	}
	n, ctx, _ := minimalHost(t)
	m := NewOperatorMessenger(ctx, n, op)
	m.SetAuthToken(op.AuthToken)

	env := m.CreateEnvelope(MsgTypePs, []byte("test"))
	if len(env.Token) != 32 {
		t.Fatalf("token not set: got %d bytes", len(env.Token))
	}

	// Simulate implant-side check: wrong token should be rejected
	wrongToken := make([]byte, 32)
	for i := range wrongToken {
		wrongToken[i] = 0xFF
	}
	if bytes.Equal(env.Token, wrongToken) {
		t.Fatal("tampered token accidentally matches")
	}
	if bytes.Equal(env.Token, op.AuthToken) {
		// expected — token is correct. Now verify wrong token fails
		env.Token = wrongToken
		if bytes.Equal(env.Token, op.AuthToken) {
			t.Fatal("failed to set wrong token")
		}
	}
}

func TestEnvelopeSignAndVerify(t *testing.T) {
	op, err := cryptography.GenerateOperatorKey()
	if err != nil {
		t.Fatalf("generate operator key: %v", err)
	}
	n, ctx, _ := minimalHost(t)
	m := NewOperatorMessenger(ctx, n, op)

	data := []byte("necropolis command data")
	env := m.CreateEnvelope(MsgTypeExecute, data)

	signingData, err := EnvelopeSigningBytes(env)
	if err != nil {
		t.Fatalf("signing bytes: %v", err)
	}
	sig, err := op.PrivateKey.Sign(signingData)
	if err != nil {
		t.Fatalf("sign: %v", err)
	}
	env.Signature = sig

	ok, err := cryptography.Verify(op.PublicKey, signingData, sig)
	if err != nil {
		t.Fatalf("verify: %v", err)
	}
	if !ok {
		t.Fatal("signature verification failed")
	}

	// Tamper with data, re-signing bytes change, verify should fail
	data[0] ^= 0xFF
	env.Data = data
	signingData2, _ := EnvelopeSigningBytes(env)
	ok2, _ := cryptography.Verify(op.PublicKey, signingData2, sig)
	if ok2 {
		t.Fatal("tampered data passed signature verification")
	}
}

func TestTopicFormats(t *testing.T) {
	op, err := cryptography.GenerateOperatorKey()
	if err != nil {
		t.Fatalf("generate operator key: %v", err)
	}
	n, ctx, _ := minimalHost(t)
	m := NewOperatorMessenger(ctx, n, op)
	m.SetOperatorID(op.PeerID)
	opID := op.PeerID.String()

	if m.BeaconTopic() != "/b/"+opID+"/bx" {
		t.Fatalf("beacon topic: got %s", m.BeaconTopic())
	}
	if m.CommandTopic() != "/c/"+opID+"/cx" {
		t.Fatalf("command topic: got %s", m.CommandTopic())
	}
	taskTopic := m.TaskTopic("test-implant-id")
	if !strings.HasPrefix(taskTopic, "/b/") || !strings.HasSuffix(taskTopic, "/tx/test-implant-id") {
		t.Fatalf("task topic: got %s", taskTopic)
	}
}

func TestBoxEncryptDecryptViaEnvelope(t *testing.T) {
	op, err := cryptography.GenerateOperatorKey()
	if err != nil {
		t.Fatalf("generate operator key: %v", err)
	}
	n, ctx, _ := minimalHost(t)
	m := NewOperatorMessenger(ctx, n, op)

	plain := []byte("secret implant command")
	env := m.CreateEnvelope(MsgTypeExecute, plain)

	// Encrypt the data field
	if op.BoxKeys == nil {
		t.Skip("no box keys")
	}
	encrypted, err := cryptography.EncryptMessage(env.Data, op.BoxKeys.PublicKey)
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}
	if bytes.Equal(encrypted, plain) {
		t.Fatal("encrypted data equals plaintext")
	}
	env.Data = encrypted

	// Decrypt
	decrypted, err := cryptography.DecryptMessage(env.Data, op.BoxKeys)
	if err != nil {
		t.Fatalf("decrypt: %v", err)
	}
	if !bytes.Equal(decrypted, plain) {
		t.Fatalf("decrypted != plain: got %q, want %q", decrypted, plain)
	}
}