Fix some review issues

Makefile

@@ -5,7 +5,7 @@ TARGET := image-full
 CONTAINER_NAME := neo_env
 
 NEED_COMMANDS := curl wget dig nc file nslookup ifconfig python3 pip3 vim
-NEED_PACKAGES := pymongo pymysql psycopg2 redis z3 secrets checklib requests pwn numpy bs4 hashpumpy dnslib regex lxml gmpy2 sympy grequests
+NEED_PACKAGES := pymongo pymysql psycopg2 redis z3 secrets checklib requests pwn numpy bs4 hashpumpy dnslib regex lxml gmpy2 sympy grequests websocket
 
 .PHONY: lint-go
 lint-go:
@@ -18,6 +18,10 @@ lint-proto:
 .PHONY: lint
 lint: lint-go lint-proto
 
+.PHONY: goimports
+goimports:
+	goimports -local github.com/c4t-but-s4d/neo -w -d $(find . -type f -name '*.go' -not -path "./proto/*")
+
 .PHONY: test
 test:
 	go test -race -timeout 1m ./...

README.md

@@ -16,10 +16,8 @@ can be a pain and wastes time. **Neo** helps in two primary ways:
    automatically.
 
 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.
+   distribution manually, 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
 

client_env/Dockerfile

@@ -6,7 +6,7 @@ WORKDIR /app
 COPY go.* ./
 COPY cmd cmd
 COPY internal internal
-COPY lib lib
+COPY proto/go proto/go
 COPY pkg pkg
 RUN --mount=type=cache,target=/root/.cache/go-build \
     --mount=type=cache,target=/go/pkg/mod \

cmd/client/cli/add.go

@@ -93,15 +93,8 @@ func (ac *addCLI) Run(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to get config from server: %w", err)
 	}
-	exists := false
-	for _, v := range state.Exploits {
-		if v.ExploitId == ac.exploitID {
-			exists = true
-			break
-		}
-	}
 
-	if exists {
+	if getExploitFromState(state, ac.exploitID) != nil {
 		fmt.Println("The exploit with this id already exists. Do you want to override (add new version) y/N ?")
 		var tmp string
 		if _, err := fmt.Scanln(&tmp); err != nil {

cmd/client/cli/cli.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"sync"
 
 	"github.com/c4t-but-s4d/neo/internal/client"
 	"github.com/c4t-but-s4d/neo/pkg/grpcauth"
@@ -25,19 +24,16 @@ type NeoCLI interface {
 type baseCLI struct {
 	cfg *client.Config
 
-	mu       sync.Mutex
 	clientID string
 }
 
 func (cmd *baseCLI) client() (*client.Client, error) {
-	var opts []grpc.DialOption
-	opts = append(
-		opts,
+	opts := []grpc.DialOption{
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithDefaultCallOptions(
 			grpc.UseCompressor(gzip.Name),
 		),
-	)
+	}
 	if cmd.cfg.GrpcAuthKey != "" {
 		interceptor := grpcauth.NewClientInterceptor(cmd.cfg.GrpcAuthKey)
 		opts = append(
@@ -54,8 +50,6 @@ func (cmd *baseCLI) client() (*client.Client, error) {
 }
 
 func (cmd *baseCLI) ClientID() string {
-	cmd.mu.Lock()
-	defer cmd.mu.Unlock()
 	if cmd.clientID != "" {
 		return cmd.clientID
 	}

cmd/client/cli/dry_run.go

@@ -55,14 +55,7 @@ func (rc *dryRunCLI) Run(ctx context.Context) error {
 		return fmt.Errorf("failed to parse config: %w", err)
 	}
 
-	exists := false
-	for _, v := range state.Exploits {
-		if v.ExploitId == rc.exploitID {
-			exists = true
-			break
-		}
-	}
-	if !exists {
+	if getExploitFromState(state, rc.exploitID) == nil {
 		return fmt.Errorf("exploit %s does not exist, add it first", rc.exploitID)
 	}
 

cmd/client/cli/set_disabled.go

@@ -6,8 +6,6 @@ import (
 
 	"github.com/spf13/cobra"
 
-	epb "github.com/c4t-but-s4d/neo/proto/go/exploits"
-
 	"github.com/c4t-but-s4d/neo/internal/client"
 )
 
@@ -35,18 +33,9 @@ func (sc *setDisabledCli) Run(ctx context.Context) error {
 		return fmt.Errorf("failed to get config from server: %w", err)
 	}
 
-	var spl *epb.ExploitState
-	for _, v := range state.Exploits {
-		if v.ExploitId == sc.exploitID {
-			spl = v
-			break
-		}
-	}
-
-	if spl == nil {
+	if spl := getExploitFromState(state, sc.exploitID); spl == nil {
 		return fmt.Errorf("exploit %s does not exist", sc.exploitID)
-	}
-	if err := c.SetExploitDisabled(ctx, spl.ExploitId, sc.disabled); err != nil {
+	} else if err := c.SetExploitDisabled(ctx, spl.ExploitId, sc.disabled); err != nil {
 		return fmt.Errorf("set disabled failed: %w", err)
 	}
 

cmd/client/cli/single.go

@@ -31,15 +31,8 @@ func (sc *singleRunCLI) Run(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to get config from server: %w", err)
 	}
-	exists := false
-	for _, v := range state.Exploits {
-		if v.ExploitId == sc.exploitID {
-			exists = true
-			break
-		}
-	}
 
-	if !exists {
+	if getExploitFromState(state, sc.exploitID) == nil {
 		logrus.Fatalf("Exploit %s does not exist. Please, add it first.", sc.exploitID)
 	}
 	if err := c.SingleRun(ctx, sc.exploitID); err != nil {

cmd/client/cli/utils.go

@@ -1,7 +1,23 @@
 package cli
 
-import "bytes"
+import (
+	"bytes"
+
+	"github.com/samber/lo"
+
+	"github.com/c4t-but-s4d/neo/proto/go/exploits"
+)
 
 func isBinary(data []byte) bool {
 	return bytes.Equal(data[:4], []byte("\x7fELF"))
 }
+
+func getExploitFromState(state *exploits.ServerState, exploitID string) *exploits.ExploitState {
+	exp, ok := lo.Find(state.Exploits, func(s *exploits.ExploitState) bool {
+		return s.ExploitId == exploitID
+	})
+	if !ok {
+		return nil
+	}
+	return exp
+}

internal/client/client.go

@@ -7,7 +7,6 @@ import (
 	"io"
 
 	"github.com/sirupsen/logrus"
-	"google.golang.org/protobuf/types/known/emptypb"
 
 	epb "github.com/c4t-but-s4d/neo/proto/go/exploits"
 
@@ -41,10 +40,9 @@ func (nc *Client) GetServerState(ctx context.Context) (*epb.ServerState, error)
 	resp, err := nc.exploits.Ping(
 		ctx,
 		&epb.PingRequest{
+			ClientId: nc.ID,
 			Payload: &epb.PingRequest_ServerInfoRequest{
-				ServerInfoRequest: &epb.PingRequest_ServerInfo{
-					ClientId: nc.ID,
-				},
+				ServerInfoRequest: &epb.PingRequest_ServerInfo{},
 			},
 		},
 	)
@@ -58,10 +56,10 @@ func (nc *Client) Heartbeat(ctx context.Context) (*epb.ServerState, error) {
 	resp, err := nc.exploits.Ping(
 		ctx,
 		&epb.PingRequest{
+			ClientId: nc.ID,
 			Payload: &epb.PingRequest_HeartbeatRequest{
 				HeartbeatRequest: &epb.PingRequest_Heartbeat{
-					ClientId: nc.ID,
-					Weight:   int32(nc.Weight),
+					Weight: int32(nc.Weight),
 				},
 			},
 		},
@@ -76,10 +74,9 @@ func (nc *Client) Leave(ctx context.Context) error {
 	if _, err := nc.exploits.Ping(
 		ctx,
 		&epb.PingRequest{
+			ClientId: nc.ID,
 			Payload: &epb.PingRequest_LeaveRequest{
-				LeaveRequest: &epb.PingRequest_Leave{
-					ClientId: nc.ID,
-				},
+				LeaveRequest: &epb.PingRequest_Leave{},
 			},
 		},
 	); err != nil {
@@ -138,7 +135,7 @@ func (nc *Client) UploadFile(ctx context.Context, r io.Reader) (*fspb.FileInfo,
 }
 
 func (nc *Client) BroadcastCommand(ctx context.Context, command string) error {
-	req := &epb.Command{Command: command}
+	req := &epb.BroadcastRequest{Command: command}
 	if _, err := nc.exploits.BroadcastCommand(ctx, req); err != nil {
 		return fmt.Errorf("making broadcast command request: %w", err)
 	}
@@ -168,13 +165,13 @@ func (nc *Client) SetExploitDisabled(ctx context.Context, id string, disabled bo
 	return nil
 }
 
-func (nc *Client) ListenBroadcasts(ctx context.Context) (<-chan *epb.Command, error) {
-	stream, err := nc.exploits.BroadcastRequests(ctx, &emptypb.Empty{})
+func (nc *Client) ListenBroadcasts(ctx context.Context) (<-chan *epb.BroadcastSubscribeResponse, error) {
+	stream, err := nc.exploits.BroadcastSubscribe(ctx, &epb.BroadcastSubscribeRequest{})
 	if err != nil {
 		return nil, fmt.Errorf("creating broadcast requests stream: %w", err)
 	}
 
-	results := make(chan *epb.Command)
+	results := make(chan *epb.BroadcastSubscribeResponse)
 	go func() {
 		defer close(results)
 		for {
@@ -194,13 +191,13 @@ func (nc *Client) ListenBroadcasts(ctx context.Context) (<-chan *epb.Command, er
 	return results, nil
 }
 
-func (nc *Client) ListenSingleRuns(ctx context.Context) (<-chan *epb.SingleRunRequest, error) {
-	stream, err := nc.exploits.SingleRunRequests(ctx, &emptypb.Empty{})
+func (nc *Client) ListenSingleRuns(ctx context.Context) (<-chan *epb.SingleRunSubscribeResponse, error) {
+	stream, err := nc.exploits.SingleRunSubscribe(ctx, &epb.SingleRunSubscribeRequest{})
 	if err != nil {
 		return nil, fmt.Errorf("creating single run requests stream: %w", err)
 	}
 
-	results := make(chan *epb.SingleRunRequest)
+	results := make(chan *epb.SingleRunSubscribeResponse)
 	go func() {
 		defer close(results)
 		for {

internal/exploit/runner.go

@@ -42,7 +42,7 @@ func NewRunner(
 		client:         c,
 		maxJobs:        maxJobs,
 		maxEndlessJobs: maxEndlessJobs,
-		singleRuns:     make(chan *epb.SingleRunRequest),
+		singleRuns:     make(chan *epb.SingleRunSubscribeResponse),
 		restarts:       make(chan struct{}, 1),
 		logSender:      logSender,
 		metricsPusher: push.
@@ -66,7 +66,7 @@ type Runner struct {
 	maxJobs        int
 	maxEndlessJobs int
 
-	singleRuns chan *epb.SingleRunRequest
+	singleRuns chan *epb.SingleRunSubscribeResponse
 	restarts   chan struct{}
 
 	logger *logrus.Entry
@@ -344,7 +344,7 @@ func (r *Runner) startPushingMetrics(ctx context.Context) {
 	}
 }
 
-func (r *Runner) handleBroadcastCommand(ctx context.Context, cmd *epb.Command) error {
+func (r *Runner) handleBroadcastCommand(ctx context.Context, cmd *epb.BroadcastSubscribeResponse) error {
 	c := exec.CommandContext(ctx, "/bin/bash", "-c", cmd.Command)
 	c.Stdout = os.Stdout
 	c.Stderr = os.Stderr