Files

120 lines
3.2 KiB
Go

// 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)
}