Multiple changes:

- Add example exploits for testing
- Running client in network mode host
- Major refactoring in tests & proto
- Fix distribution tests
- Better wording in readme
- Pubsub rewritten in generics
- More tests

README.md

@@ -8,15 +8,18 @@ Client + server for exploit distribution during Attack & Defence CTF competition
 
 ## Why
 
-Usually during large CTFs, where the regular laptop can't run the exploits on all participants, teams start some cloud
-servers to run these exploits on, but the administration and access management can be a pain and wastes time. **Neo**
-can solve this problem by doing 2 things:
+Usually during large CTFs, where the regular laptop can't run the exploits on all participants, teams rent cloud
+servers to run the exploits. However, uploading, managing and monitoring the exploits on a remote machine
+can be a pain and wastes time. **Neo** helps in two primary ways:
 
-1. Every player can start an instance of Neo client and attack a proportional part of the whole team pool automatically.
+1. Every player can start an instance of **Neo client** and attack a proportional part of the whole team pool
+   automatically.
 
-2. Exploit writers don't copy the newly-created exploits to the exploit server, or manage the distribution by hand, but
-   rather submit them to the **Neo server** using the same client, and the server does all the work, running the exploit
-   on available clients.
+2. Exploit writers don't upload the newly-created exploits to the exploit server, neither do they manage the
+   distribution
+   by hand, but rather submit them to the **Neo server** using the same client, and the server does all the work,
+   distributing the exploit
+   among the available clients.
 
 ## Usage
 
@@ -33,13 +36,14 @@ Neo uses the exploit farm to acquire the team list and submit the flags. The pro
 - `POST /api/post_flags` will receive an array of mappings with keys `flag`, `sploit` and `team`. `sploit` is the
   exploit name for statistics, and `team` is the team name.
 
-Farm password will be passed in `Authorization` and `X-Token` headers, so the protocol is compatible with **
-DestructiveFarm**.
+Farm password will be passed in `Authorization` and `X-Token` headers, so the protocol is compatible with
+**DestructiveFarm**.
 
 ### Server
 
-Server coordinates the clients and distributes teams to attack among them. It must have access to the farm, and all
-clients must have access to the server, so you might want to start it somewhere with the public IP address.
+Server coordinates the clients and distributes targets among them. It must have access to the farm, and all
+clients must have access both to the server and the farm, so you might want to start it somewhere with the public IP
+address.
 
 To start the server:
 

client_env/start.sh

@@ -31,6 +31,7 @@ else
     --security-opt apparmor=unconfined \
     --cap-add=NET_ADMIN \
     --privileged \
+    --network host \
     --name "${CONTAINER_NAME}" \
     --hostname "${CONTAINER_NAME}" \
     "${IMAGE}" \

cmd/client/cli/add.go

@@ -18,6 +18,7 @@ import (
 
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
+	"google.golang.org/protobuf/types/known/durationpb"
 
 	neopb "neo/lib/genproto/neo"
 )
@@ -93,8 +94,8 @@ func (ac *addCLI) Run(ctx context.Context) error {
 		return fmt.Errorf("failed to get config from server: %w", err)
 	}
 	exists := false
-	for _, v := range state.GetExploits() {
-		if v.GetExploitId() == ac.exploitID {
+	for _, v := range state.Exploits {
+		if v.ExploitId == ac.exploitID {
 			exists = true
 			break
 		}
@@ -150,8 +151,8 @@ func (ac *addCLI) Run(ctx context.Context) error {
 		Config: &neopb.ExploitConfiguration{
 			Entrypoint: file,
 			IsArchive:  ac.isArchive,
-			RunEvery:   ac.runEvery.String(),
-			Timeout:    ac.timeout.String(),
+			RunEvery:   durationpb.New(ac.runEvery),
+			Timeout:    durationpb.New(ac.timeout),
 		},
 		Endless:  ac.endless,
 		Disabled: ac.disabled,

cmd/client/cli/dry_run.go

@@ -51,14 +51,14 @@ func (rc *dryRunCLI) Run(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to get config from server: %w", err)
 	}
-	cfg, err := config.FromProto(state.GetConfig())
+	cfg, err := config.FromProto(state.Config)
 	if err != nil {
 		return fmt.Errorf("failed to parse config: %w", err)
 	}
 
 	exists := false
-	for _, v := range state.GetExploits() {
-		if v.GetExploitId() == rc.exploitID {
+	for _, v := range state.Exploits {
+		if v.ExploitId == rc.exploitID {
 			exists = true
 			break
 		}
@@ -68,15 +68,15 @@ func (rc *dryRunCLI) Run(ctx context.Context) error {
 	}
 
 	storage := exploit.NewStorage(exploit.NewCache(), rc.baseCLI.c.ExploitDir, c)
-	storage.UpdateExploits(ctx, state.GetExploits())
+	storage.UpdateExploits(ctx, state.Exploits)
 	ex, ok := storage.Exploit(rc.exploitID)
 	if !ok {
 		return fmt.Errorf("failed to find exploit '%s' in storage", rc.exploitID)
 	}
 
 	allTeams := make(map[string]string)
-	for _, tbuck := range state.GetClientTeamMap() {
-		for k, v := range tbuck.GetTeams() {
+	for _, tbuck := range state.ClientTeamMap {
+		for k, v := range tbuck.Teams {
 			allTeams[k] = v
 		}
 	}

cmd/client/cli/info.go

@@ -29,19 +29,19 @@ func (ic *infoCLI) Run(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("making ping config request: %w", err)
 	}
-	cfg, err := config.FromProto(state.GetConfig())
+	cfg, err := config.FromProto(state.Config)
 	if err != nil {
 		return fmt.Errorf("unmarshalling config: %w", err)
 	}
 	fmt.Printf("config: %+v\n", cfg)
 	fmt.Println("IPs buckets: ")
-	for k, v := range state.GetClientTeamMap() {
+	for k, v := range state.ClientTeamMap {
 		fmt.Print(k, ": [")
-		fmt.Printf("%+v", v.GetTeams())
+		fmt.Printf("%+v", v.Teams)
 		fmt.Println("]")
 	}
 	fmt.Println("Exploits: ")
-	for _, e := range state.GetExploits() {
+	for _, e := range state.Exploits {
 		fmt.Printf("%+v\n", e)
 	}
 	return nil

cmd/client/cli/set_disabled.go

@@ -36,8 +36,8 @@ func (sc *setDisabledCli) Run(ctx context.Context) error {
 	}
 
 	var spl *neopb.ExploitState
-	for _, v := range state.GetExploits() {
-		if v.GetExploitId() == sc.exploitID {
+	for _, v := range state.Exploits {
+		if v.ExploitId == sc.exploitID {
 			spl = v
 			break
 		}
@@ -46,7 +46,7 @@ func (sc *setDisabledCli) Run(ctx context.Context) error {
 	if spl == nil {
 		return fmt.Errorf("exploit %s does not exist", sc.exploitID)
 	}
-	if err := c.SetExploitDisabled(ctx, spl.GetExploitId(), sc.disabled); err != nil {
+	if err := c.SetExploitDisabled(ctx, spl.ExploitId, sc.disabled); err != nil {
 		return fmt.Errorf("set disabled failed: %w", err)
 	}
 

cmd/client/cli/single.go

@@ -34,8 +34,8 @@ func (sc *singleRunCLI) Run(ctx context.Context) error {
 		return fmt.Errorf("failed to get config from server: %w", err)
 	}
 	exists := false
-	for _, v := range state.GetExploits() {
-		if v.GetExploitId() == sc.exploitID {
+	for _, v := range state.Exploits {
+		if v.ExploitId == sc.exploitID {
 			exists = true
 			break
 		}

examples/simple.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+
+import base64
+import os
+import time
+
+for i in range(5):
+    print(base64.b64encode(os.urandom(23), altchars=b'AB').decode().upper())
+    time.sleep(1)

examples/spammy.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python
+
+import base64
+import os
+
+for i in range(1000):
+    print('A' * 2 * 1024 * 1024)
+    print(base64.b64encode(os.urandom(23), altchars=b'AB').decode().upper())

examples/timeout.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python
+
+import base64
+import os
+import time
+
+while True:
+    print(base64.b64encode(os.urandom(23), altchars=b'AB').decode().upper())
+    time.sleep(1)

examples/unkillable.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python
+
+import base64
+import os
+import time
+
+while True:
+    try:
+        print(base64.b64encode(os.urandom(23), altchars=b'AB').decode().upper())
+        time.sleep(1)
+    except BaseException:  # noqa
+        pass

internal/client/client.go

@@ -7,6 +7,7 @@ import (
 	"io"
 
 	"github.com/sirupsen/logrus"
+	"google.golang.org/protobuf/types/known/emptypb"
 
 	"neo/pkg/filestream"
 
@@ -38,7 +39,7 @@ func (nc *Client) Ping(ctx context.Context, t neopb.PingRequest_PingType) (*neop
 	if err != nil {
 		return nil, fmt.Errorf("making ping request: %w", err)
 	}
-	return resp.GetState(), nil
+	return resp.State, nil
 }
 
 func (nc *Client) Exploit(ctx context.Context, id string) (*neopb.ExploitResponse, error) {
@@ -58,7 +59,7 @@ func (nc *Client) UpdateExploit(ctx context.Context, state *neopb.ExploitState)
 	if err != nil {
 		return nil, fmt.Errorf("aking update exploit request: %w", err)
 	}
-	return resp.GetState(), nil
+	return resp.State, nil
 }
 
 func (nc *Client) DownloadFile(ctx context.Context, info *neopb.FileInfo, out io.Writer) error {
@@ -112,7 +113,7 @@ func (nc *Client) SetExploitDisabled(ctx context.Context, id string, disabled bo
 		return fmt.Errorf("fetching current exploit config: %w", err)
 	}
 
-	req := &neopb.UpdateExploitRequest{State: resp.GetState()}
+	req := &neopb.UpdateExploitRequest{State: resp.State}
 	req.State.Disabled = disabled
 
 	if _, err := nc.c.UpdateExploit(ctx, req); err != nil {
@@ -122,7 +123,7 @@ func (nc *Client) SetExploitDisabled(ctx context.Context, id string, disabled bo
 }
 
 func (nc *Client) ListenBroadcasts(ctx context.Context) (<-chan *neopb.Command, error) {
-	stream, err := nc.c.BroadcastRequests(ctx, &neopb.Empty{})
+	stream, err := nc.c.BroadcastRequests(ctx, &emptypb.Empty{})
 	if err != nil {
 		return nil, fmt.Errorf("creating broadcast requests stream: %w", err)
 	}
@@ -148,7 +149,7 @@ func (nc *Client) ListenBroadcasts(ctx context.Context) (<-chan *neopb.Command,
 }
 
 func (nc *Client) ListenSingleRuns(ctx context.Context) (<-chan *neopb.SingleRunRequest, error) {
-	stream, err := nc.c.SingleRunRequests(ctx, &neopb.Empty{})
+	stream, err := nc.c.SingleRunRequests(ctx, &emptypb.Empty{})
 	if err != nil {
 		return nil, fmt.Errorf("creating single run requests stream: %w", err)
 	}

internal/config/config.go

@@ -6,6 +6,8 @@ import (
 	"time"
 
 	neopb "neo/lib/genproto/neo"
+
+	"google.golang.org/protobuf/types/known/durationpb"
 )
 
 type Config struct {
@@ -22,8 +24,8 @@ func ToProto(c *Config) *neopb.Config {
 		FarmUrl:      c.FarmURL,
 		FarmPassword: c.FarmPassword,
 		FlagRegexp:   c.FlagRegexp.String(),
-		PingEvery:    c.PingEvery.String(),
-		SubmitEvery:  c.SubmitEvery.String(),
+		PingEvery:    durationpb.New(c.PingEvery),
+		SubmitEvery:  durationpb.New(c.SubmitEvery),
 		Environ:      c.Environ,
 	}
 }
@@ -33,17 +35,13 @@ func FromProto(config *neopb.Config) (*Config, error) {
 		cfg Config
 		err error
 	)
-	if cfg.FlagRegexp, err = regexp.Compile(config.GetFlagRegexp()); err != nil {
+	if cfg.FlagRegexp, err = regexp.Compile(config.FlagRegexp); err != nil {
 		return nil, fmt.Errorf("compiling regex: %w", err)
 	}
-	if cfg.PingEvery, err = time.ParseDuration(config.GetPingEvery()); err != nil {
-		return nil, fmt.Errorf("parsing ping interval: %w", err)
-	}
-	if cfg.SubmitEvery, err = time.ParseDuration(config.GetSubmitEvery()); err != nil {
-		return nil, fmt.Errorf("parsing submit interval: %w", err)
-	}
-	cfg.FarmURL = config.GetFarmUrl()
-	cfg.FarmPassword = config.GetFarmPassword()
-	cfg.Environ = config.GetEnviron()
+	cfg.FarmURL = config.FarmUrl
+	cfg.FarmPassword = config.FarmPassword
+	cfg.PingEvery = config.PingEvery.AsDuration()
+	cfg.SubmitEvery = config.SubmitEvery.AsDuration()
+	cfg.Environ = config.Environ
 	return &cfg, nil
 }

internal/exploit/cache.go

@@ -42,7 +42,7 @@ func (c *Cache) Diff(exs []*neopb.ExploitState) (diff []*neopb.ExploitState, res
 	c.m.RLock()
 	defer c.m.RUnlock()
 	for _, ex := range exs {
-		state, exists := c.states[ex.GetExploitId()]
+		state, exists := c.states[ex.ExploitId]
 		if !exists {
 			if ex.Endless {
 				restartEndless = true
@@ -53,7 +53,7 @@ func (c *Cache) Diff(exs []*neopb.ExploitState) (diff []*neopb.ExploitState, res
 		if state.Endless != ex.Endless {
 			restartEndless = true
 		}
-		if state.Version != ex.GetVersion() {
+		if state.Version != ex.Version {
 			diff = append(diff, ex)
 		}
 	}

internal/exploit/cache_test.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/testing/protocmp"
 
 	neopb "neo/lib/genproto/neo"
@@ -195,9 +196,7 @@ func TestCache_Diff(t *testing.T) {
 		if diff := cmp.Diff(tc.diff, got, protocmp.Transform(), cmpopts.EquateEmpty()); diff != "" {
 			t.Errorf("Diff() mismatch (-want +got):\n%s", diff)
 		}
-		if re != tc.re {
-			t.Errorf("Diff(): expected restartEndless = %t, got = %t", tc.re, re)
-		}
+		require.Equal(t, tc.re, re)
 	}
 }
 

internal/exploit/runner.go

@@ -62,7 +62,7 @@ func (r *Runner) Run(ctx context.Context) error {
 	}
 
 	r.printGreeting()
-	r.updateTeams(state.GetClientTeamMap())
+	r.updateTeams(state.ClientTeamMap)
 
 	wg := new(sync.WaitGroup)
 
@@ -106,12 +106,12 @@ func (r *Runner) pingHeartbeat(ctx context.Context) (*neopb.ServerState, error)
 	if err != nil {
 		return nil, fmt.Errorf("sending ping request: %w", err)
 	}
-	cfg, err := config.FromProto(state.GetConfig())
+	cfg, err := config.FromProto(state.Config)
 	if err != nil {
 		return nil, fmt.Errorf("parsing config: %w", err)
 	}
 	r.updateConfig(cfg)
-	r.updateExploits(ctx, state.GetExploits())
+	r.updateExploits(ctx, state.Exploits)
 	return state, nil
 }
 
@@ -144,7 +144,7 @@ func (r *Runner) eventLoop(ctx context.Context) {
 			if err != nil {
 				logrus.Errorf("Error sending recurrent heartbeat: %v", err)
 			}
-			if r.updateTeams(state.GetClientTeamMap()) {
+			if r.updateTeams(state.ClientTeamMap) {
 				r.restartSimple()
 				r.restartEndless()
 			}
@@ -306,7 +306,7 @@ func (r *Runner) updateTeams(buckets map[string]*neopb.TeamBucket) bool {
 		logrus.Errorf("Failed to find IPs in state for client: %s", cid)
 		return false
 	}
-	teams := ipbuck.GetTeams()
+	teams := ipbuck.Teams
 
 	r.teamsLock.Lock()
 	defer r.teamsLock.Unlock()