Adding supervisor mode

Dockerfile

@@ -2,17 +2,17 @@ FROM --platform=linux/amd64 golang:1.25 AS builder
 
 RUN apt-get update && apt-get install -y git libpcap-dev
 
-# Tools dependencies
+# Tools dependencies with optimization flags
 # dnsx
-RUN go install -v github.com/projectdiscovery/dnsx/cmd/dnsx@latest
+RUN go install -ldflags="-s -w" github.com/projectdiscovery/dnsx/cmd/dnsx@latest
 # naabu
-RUN go install -v github.com/projectdiscovery/naabu/v2/cmd/naabu@latest
+RUN go install -ldflags="-s -w" github.com/projectdiscovery/naabu/v2/cmd/naabu@latest
 # httpx
-RUN go install -v github.com/projectdiscovery/httpx/cmd/httpx@latest
+RUN go install -ldflags="-s -w" github.com/projectdiscovery/httpx/cmd/httpx@latest
 # tlsx
-RUN go install -v github.com/projectdiscovery/tlsx/cmd/tlsx@latest
+RUN go install -ldflags="-s -w" github.com/projectdiscovery/tlsx/cmd/tlsx@latest
 # nuclei
-RUN go install -v github.com/projectdiscovery/nuclei/v3/cmd/nuclei@latest
+RUN go install -ldflags="-s -w" github.com/projectdiscovery/nuclei/v3/cmd/nuclei@latest
 
 
 # Copy source code
@@ -25,22 +25,28 @@ COPY . .
 # CGO_ENABLED=1 is required for libpcap/gopacket support (passive discovery feature)
 RUN CGO_ENABLED=1 GOOS=linux go build -ldflags="-s -w" -o /go/bin/pd-agent ./cmd/pd-agent/main.go
 
-FROM --platform=linux/amd64 ubuntu:latest
+# Clean Go module cache to reduce image size
+RUN go clean -modcache && \
+    rm -rf /root/.cache/go-build
+
+FROM --platform=linux/amd64 ubuntu:22.04
 # install dependencies
 # required: libpcap-dev, chrome
-RUN apt update && apt install -y \
-    bind9-dnsutils \
-    ca-certificates \
-    nmap \
-    libpcap-dev \
-    wget \
-    gnupg \
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        ca-certificates \
+        libpcap-dev \
+        wget \
+        gnupg \
     && wget -q -O - https://dl.google.com/linux/linux_signing_key.pub | gpg --dearmor -o /usr/share/keyrings/google-chrome.gpg \
     && echo "deb [arch=amd64 signed-by=/usr/share/keyrings/google-chrome.gpg] http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list \
-    && apt update \
-    && apt install -y google-chrome-stable \
-    && apt clean \
-    && rm -rf /var/lib/apt/lists/*
+    && apt-get update \
+    && apt-get install -y --no-install-recommends google-chrome-stable \
+    && apt-get purge -y wget gnupg \
+    && apt-get autoremove -y \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/* \
+    && rm -rf /usr/share/doc /usr/share/man /usr/share/locale /usr/share/info
 
 # Set environment variables for Chrome
 ENV CHROME_BIN=/usr/bin/google-chrome-stable
@@ -57,8 +63,9 @@ COPY --from=builder /go/bin/nuclei /usr/local/bin/
 # Copy agent binary
 COPY --from=builder /go/bin/pd-agent /usr/local/bin/pd-agent
 
-# Create writable output directory for existing ubuntu user (UID 1000)
-RUN mkdir -p /home/ubuntu/output && \
+# Create ubuntu user and writable output directory
+RUN useradd -m -u 1000 ubuntu && \
+    mkdir -p /home/ubuntu/output && \
     chown -R ubuntu:ubuntu /home/ubuntu
 
 # Set default environment variables (can be overridden at runtime)

README.md

@@ -14,6 +14,7 @@
   <a href="#pd-agent">PD Agent</a> •
   <a href="#installation">Installation</a> •
   <a href="#quick-start">Quick Start</a> •
+  <a href="#supervisor-mode">Supervisor Mode</a> •
   <a href="#system-installation">System Installation</a> •
   <a href="https://discord.gg/projectdiscovery">Join Discord</a>
 
@@ -65,6 +66,38 @@ kubectl get pods -n pd-agent -l app=pd-agent
 
 The agent automatically discovers Kubernetes cluster subnets (nodes, pods, services) for scanning. See [examples/README.md](examples/README.md) for detailed instructions and customization options.
 
+### Supervisor Mode
+
+Supervisor mode allows pd-agent to manage its own deployment in Docker or Kubernetes, automatically handling updates, restarts, and lifecycle management.
+
+#### Prerequisites
+
+- **macOS/Windows**: Docker Desktop must be installed and running
+  - On Windows, Docker Desktop has automatic integration with WSL2
+- **Linux**: Docker must be installed and running
+
+#### Usage
+
+Run pd-agent in supervisor mode with Docker (default):
+
+```bash
+pd-agent -supervisor-mode docker
+```
+
+Or use Kubernetes:
+
+```bash
+pd-agent -supervisor-mode kubernetes
+```
+
+The supervisor will:
+- Automatically pull and deploy the latest pd-agent Docker image
+- Monitor the deployment and restart if it crashes
+- Handle image updates automatically
+- Manage the container/pod lifecycle
+
+**Note**: Supervisor mode requires Docker or Kubernetes to be available and properly configured. The supervisor runs the agent in a container/pod, so all agent configuration (environment variables, flags) should be passed as normal.
+
 ### Network Discovery
 
 The agent automatically discovers local network subnets and reports them to the platform:

cmd/pd-agent/main.go

@@ -29,6 +29,7 @@ import (
 	"github.com/projectdiscovery/pd-agent/pkg"
 	"github.com/projectdiscovery/pd-agent/pkg/client"
 	"github.com/projectdiscovery/pd-agent/pkg/scanlog"
+	"github.com/projectdiscovery/pd-agent/pkg/supervisor"
 	"github.com/projectdiscovery/pd-agent/pkg/types"
 	"github.com/projectdiscovery/utils/batcher"
 	envutil "github.com/projectdiscovery/utils/env"
@@ -69,7 +70,7 @@ func getAllNucleiTemplates(templateDir string) ([]string, error) {
 var (
 	PDCPApiKey                = envutil.GetEnvOrDefault("PDCP_API_KEY", "")
 	TeamIDEnv                 = envutil.GetEnvOrDefault("PDCP_TEAM_ID", "")
-	AgentTagsEnv              = envutil.GetEnvOrDefault("PDCP_AGENT_TAGS", "default")
+	AgentTagsEnv              = envutil.GetEnvOrDefault("PDCP_AGENT_TAGS", "")
 	PdcpApiServer             = envutil.GetEnvOrDefault("PDCP_API_SERVER", "https://api.projectdiscovery.io")
 	ChunkParallelismEnv       = envutil.GetEnvOrDefault("PDCP_CHUNK_PARALLELISM", "1")
 	ScanParallelismEnv        = envutil.GetEnvOrDefault("PDCP_SCAN_PARALLELISM", "1")
@@ -90,6 +91,7 @@ type Options struct {
 	ScanParallelism        int  // Number of scans to process in parallel
 	EnumerationParallelism int  // Number of enumerations to process in parallel
 	KeepOutputFiles        bool // If true, don't delete output files after processing
+	SupervisorMode         string // Supervisor mode: "docker" or "kubernetes" (default: empty, disabled)
 }
 
 // ScanCache represents cached scan execution information
@@ -1429,6 +1431,11 @@ func (r *Runner) processChunks(ctx context.Context, taskID, taskType string, exe
 			if err == nil {
 				break
 			}
+			// If we get "no more chunks", terminate immediately without retrying
+			if err != nil && err.Error() == "no more chunks" {
+				r.logHelper("INFO", fmt.Sprintf("No more chunks available for %s ID: %s", taskType, taskID))
+				goto Complete
+			}
 			currentErr := err.Error()
 			if currentErr == lastErr {
 				// If we get the same error multiple times, likely the task is complete
@@ -1896,6 +1903,11 @@ func (r *Runner) getTaskChunk(ctx context.Context, taskID string, done bool) (*T
 		return nil, fmt.Errorf("error unmarshaling response: %v", err)
 	}
 
+	// Check if the unmarshaled struct is empty (no more chunks)
+	if taskChunk.ChunkID == "" {
+		return nil, fmt.Errorf("no more chunks")
+	}
+
 	return &taskChunk, nil
 }
 
@@ -2597,12 +2609,21 @@ func parseOptions() *Options {
 		flagSet.IntVarP(&options.ChunkParallelism, "chunk-parallelism", "c", defaultChunkParallelism, "number of chunks to process in parallel"),
 		flagSet.IntVarP(&options.ScanParallelism, "scan-parallelism", "s", defaultScanParallelism, "number of scans to process in parallel"),
 		flagSet.IntVarP(&options.EnumerationParallelism, "enumeration-parallelism", "e", defaultEnumerationParallelism, "number of enumerations to process in parallel"),
+		flagSet.StringVar(&options.SupervisorMode, "supervisor-mode", "", "run as supervisor: \"docker\" or \"kubernetes\" (default: empty, disabled)"),
 	)
 
 	if err := flagSet.Parse(); err != nil {
 		slog.Error("error", "error", err)
 	}
 
+	// Validate supervisor mode
+	if options.SupervisorMode != "" {
+		if options.SupervisorMode != "docker" && options.SupervisorMode != "kubernetes" {
+			slog.Error("Invalid supervisor mode", "mode", options.SupervisorMode, "valid", "docker or kubernetes")
+			options.SupervisorMode = "" // disable supervisor mode if invalid
+		}
+	}
+
 	// Parse environment variables (env vars take precedence as defaults)
 	if agentTags := os.Getenv("PDCP_AGENT_TAGS"); agentTags != "" && len(options.AgentTags) == 0 {
 		options.AgentTags = goflags.StringSlice(strings.Split(agentTags, ","))
@@ -2681,6 +2702,12 @@ func main() {
 
 	options := parseOptions()
 
+	// If supervisor mode is enabled, run supervisor instead of direct agent
+	if options.SupervisorMode != "" {
+		runSupervisorMode(options)
+		return
+	}
+
 	// Check prerequisites before starting the agent
 	prerequisites := pkg.CheckAllPrerequisites()
 	var missingTools []string
@@ -2718,3 +2745,44 @@ func main() {
 		os.Exit(1)
 	}
 }
+
+// runSupervisorMode runs the agent in supervisor mode
+func runSupervisorMode(options *Options) {
+	// Convert Options to supervisor.AgentOptions
+	agentOptions := &supervisor.AgentOptions{
+		TeamID:                 options.TeamID,
+		AgentID:                options.AgentId,
+		AgentTags:              []string(options.AgentTags),
+		AgentNetworks:          []string(options.AgentNetworks),
+		AgentOutput:            options.AgentOutput,
+		AgentName:              options.AgentName,
+		Verbose:                options.Verbose,
+		PassiveDiscovery:       options.PassiveDiscovery,
+		ChunkParallelism:       options.ChunkParallelism,
+		ScanParallelism:        options.ScanParallelism,
+		EnumerationParallelism: options.EnumerationParallelism,
+		KeepOutputFiles:        options.KeepOutputFiles,
+	}
+
+	// Generate agent ID if not set
+	if agentOptions.AgentID == "" {
+		agentOptions.AgentID = xid.New().String()
+	}
+
+	// Create supervisor with specified provider
+	sup, err := supervisor.NewSupervisorWithProvider(agentOptions, options.SupervisorMode)
+	if err != nil {
+		gologger.Fatal().Msgf("Could not create supervisor: %v", err)
+		os.Exit(1)
+	}
+
+	// Setup signal handlers
+	ctx := sup.SetupSignalHandlers(context.Background())
+
+	// Run supervisor
+	if err := sup.Run(ctx); err != nil {
+		gologger.Fatal().Msgf("Supervisor error: %v", err)
+		os.Exit(1)
+	}
+	gologger.Info().Msg("Supervisor terminated")
+}

go.mod

@@ -3,15 +3,21 @@ module github.com/projectdiscovery/pd-agent
 go 1.24.2
 
 require (
+	github.com/Masterminds/semver/v3 v3.4.0
+	github.com/docker/go-sdk/client v0.1.0-alpha012
 	github.com/dustin/go-humanize v1.0.1
+	github.com/google/go-github/v62 v62.0.0
 	github.com/google/gopacket v1.1.19
+	github.com/moby/moby/api v1.52.0
+	github.com/moby/moby/client v0.1.0
 	github.com/projectdiscovery/gcache v0.0.0-20241015120333-12546c6e3f4c
 	github.com/projectdiscovery/goflags v0.1.74
 	github.com/projectdiscovery/gologger v1.1.61
 	github.com/projectdiscovery/mapcidr v1.1.97
 	github.com/projectdiscovery/naabu/v2 v2.3.7
 	github.com/projectdiscovery/nuclei/v3 v3.5.1
 	github.com/projectdiscovery/utils v0.7.3
+	github.com/rhysd/go-github-selfupdate v1.2.3
 	github.com/rs/xid v1.6.0
 	github.com/shirou/gopsutil/v3 v3.24.5
 	github.com/tidwall/gjson v1.18.0
@@ -23,7 +29,7 @@ require (
 
 require (
 	aead.dev/minisign v0.2.0 // indirect
-	github.com/Masterminds/semver/v3 v3.2.1 // indirect
+	github.com/Microsoft/go-winio v0.6.2 // indirect
 	github.com/Mzack9999/gcache v0.0.0-20230410081825-519e28eab057 // indirect
 	github.com/Mzack9999/go-http-digest-auth-client v0.6.1-0.20220414142836-eb8883508809 // indirect
 	github.com/STARRY-S/zip v0.2.3 // indirect
@@ -35,11 +41,13 @@ require (
 	github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
+	github.com/blang/semver v3.5.1+incompatible // indirect
 	github.com/bodgit/plumbing v1.3.0 // indirect
 	github.com/bodgit/sevenzip v1.6.1 // indirect
 	github.com/bodgit/windows v1.0.1 // indirect
 	github.com/bytedance/sonic v1.14.0 // indirect
 	github.com/bytedance/sonic/loader v0.3.0 // indirect
+	github.com/caarlos0/env/v11 v11.3.1 // indirect
 	github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect
 	github.com/charmbracelet/glamour v0.10.0 // indirect
 	github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 // indirect
@@ -51,16 +59,24 @@ require (
 	github.com/cloudflare/circl v1.6.1 // indirect
 	github.com/cloudwego/base64x v0.1.5 // indirect
 	github.com/cnf/structhash v0.0.0-20250313080605-df4c6cc74a9a // indirect
+	github.com/containerd/errdefs v1.0.0 // indirect
+	github.com/containerd/errdefs/pkg v0.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/dimchansky/utfbom v1.1.1 // indirect
+	github.com/distribution/reference v0.6.0 // indirect
 	github.com/dlclark/regexp2 v1.11.5 // indirect
+	github.com/docker/go-connections v0.6.0 // indirect
+	github.com/docker/go-sdk/config v0.1.0-alpha012 // indirect
+	github.com/docker/go-sdk/context v0.1.0-alpha012 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/dsnet/compress v0.0.2-0.20230904184137-39efe44ab707 // indirect
 	github.com/emicklei/go-restful/v3 v3.12.2 // indirect
 	github.com/fatih/color v1.18.0 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/fxamacker/cbor/v2 v2.9.0 // indirect
 	github.com/gaissmai/bart v0.26.0 // indirect
 	github.com/go-logr/logr v1.4.3 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/go-openapi/jsonpointer v0.21.0 // indirect
 	github.com/go-openapi/jsonreference v0.21.0 // indirect
@@ -76,6 +92,7 @@ require (
 	github.com/gopacket/gopacket v1.2.0 // indirect
 	github.com/gorilla/css v1.0.1 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
+	github.com/inconshreveable/go-update v0.0.0-20160112193335-8152e7eb6ccf // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/klauspost/compress v1.18.2 // indirect
@@ -94,14 +111,18 @@ require (
 	github.com/mikelolasagasti/xz v1.0.1 // indirect
 	github.com/minio/minlz v1.0.1 // indirect
 	github.com/minio/selfupdate v0.6.1-0.20230907112617-f11e74f84ca7 // indirect
+	github.com/moby/docker-image-spec v1.3.1 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
 	github.com/muesli/reflow v0.3.0 // indirect
 	github.com/muesli/termenv v0.16.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/nwaples/rardecode/v2 v2.2.1 // indirect
+	github.com/opencontainers/go-digest v1.0.0 // indirect
+	github.com/opencontainers/image-spec v1.1.1 // indirect
 	github.com/pierrec/lz4/v4 v4.1.22 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect
 	github.com/projectdiscovery/asnmap v1.1.1 // indirect
 	github.com/projectdiscovery/blackrock v0.0.1 // indirect
@@ -127,7 +148,9 @@ require (
 	github.com/sorairolake/lzip-go v0.3.8 // indirect
 	github.com/spf13/afero v1.15.0 // indirect
 	github.com/spf13/pflag v1.0.6 // indirect
+	github.com/stretchr/testify v1.11.1 // indirect
 	github.com/syndtr/goleveldb v1.0.0 // indirect
+	github.com/tcnksm/go-gitconfig v0.1.2 // indirect
 	github.com/tidwall/btree v1.7.0 // indirect
 	github.com/tidwall/buntdb v1.3.1 // indirect
 	github.com/tidwall/grect v0.1.4 // indirect
@@ -150,6 +173,11 @@ require (
 	github.com/zmap/rc2 v0.0.0-20190804163417-abaa70531248 // indirect
 	github.com/zmap/zcrypto v0.0.0-20240512203510-0fef58d9a9db // indirect
 	go.etcd.io/bbolt v1.4.0 // indirect
+	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.62.0 // indirect
+	go.opentelemetry.io/otel v1.38.0 // indirect
+	go.opentelemetry.io/otel/metric v1.38.0 // indirect
+	go.opentelemetry.io/otel/trace v1.38.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.yaml.in/yaml/v2 v2.4.2 // indirect
 	go.yaml.in/yaml/v3 v3.0.4 // indirect