Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 52379c2434 | |||
| 48d17e0e9f | |||
| 6500fadafb | |||
| 6348c4c060 | |||
| 821ce5f991 | |||
| e2ec7e7234 | |||
| 0493076575 |
@@ -0,0 +1,30 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to the `theta-agent` daemon will be documented in this file.
|
||||
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
|
||||
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
||||
|
||||
## [v1.3.0] - 2026-08-04
|
||||
|
||||
### Fixed
|
||||
- **Silently ignore `heartbeat_ack`** — the server replies to the agent's own periodic heartbeat with `heartbeat_ack`. The agent had no case for it, so it fell through to the unknown-command handler, logged `Unknown command type: heartbeat_ack` every minute, and answered with a spurious error response. Heartbeat acks are fire-and-forget; the agent now ignores them silently.
|
||||
|
||||
## [v1.2.0] - 2026-08-03
|
||||
|
||||
### Added
|
||||
- **Protocol v1.1.0 Compliance**: Full alignment with `PROTOCOL.md` (v1.1.0) specification.
|
||||
- **Ed25519 Cryptographic Verification**: Verification of Ed25519 Base64 signatures for high-risk C2 commands (`reboot`, `service_restart`, `configure_ldap`, `arbitrary_bash`, `update_binary`).
|
||||
- **Enhanced Journal Log Fetcher (`fetch_logs`)**: Support for querying service-specific systemd logs (`service` parameter) with configurable line count (`lines` parameter).
|
||||
- **Pure Go Self-Update Engine (`update_binary`)**: Replaced shell script execution with pure Go HTTP client fetching, SHA256 verification, atomic file replacement, and clean daemon restart.
|
||||
|
||||
### Fixed
|
||||
- **Goroutine Leak Prevention**: Added `stopCh` lifecycle management to terminate background telemetry and heartbeat tickers upon WebSocket disconnection.
|
||||
- **Dynamic Config Rerenders**: Resolved data races when toggling capabilities or reloading `agent.yml` via `reload_config`.
|
||||
- **Config Path Uniformity**: Standardized canonical configuration file path across codebase, installer, and documentation to `/etc/theta42/agent.yml`.
|
||||
|
||||
## [v0.1.0] - 2026-08-01
|
||||
|
||||
### Added
|
||||
- Initial release of `theta-agent` Go daemon replacing legacy bash metric scripts.
|
||||
- Persistent outbound WebSocket telemetry and local capability matrix enforcement.
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
# Multi-stage build for minimal runtime image
|
||||
FROM golang:1.23-alpine AS builder
|
||||
|
||||
# Install build dependencies
|
||||
RUN apk add --no-cache git
|
||||
|
||||
WORKDIR /app
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
|
||||
COPY . .
|
||||
RUN CGO_ENABLED=0 GOOS=linux go build -v -o theta-agent .
|
||||
|
||||
# Final runtime image
|
||||
FROM alpine:latest
|
||||
RUN apk add --no-cache ca-certificates
|
||||
|
||||
WORKDIR /
|
||||
|
||||
# Create config directory
|
||||
RUN mkdir -p /etc/theta42
|
||||
|
||||
COPY --from=builder /app/theta-agent /usr/local/bin/theta-agent
|
||||
|
||||
# Run as root for system management
|
||||
USER root
|
||||
|
||||
ENTRYPOINT ["/usr/local/bin/theta-agent"]
|
||||
+93
@@ -0,0 +1,93 @@
|
||||
# Installation Guide: Theta Agent
|
||||
|
||||
Theta Agent is designed for rapid deployment across the fleet. The recommended method is via the "One-Liner" install, which marries the agent to a specific SSO Manager instance.
|
||||
|
||||
## Quick Start (The One-Liner)
|
||||
|
||||
The SSO Manager provides a pre-generated installation command. Copy and paste it into your terminal as root:
|
||||
|
||||
### Option A: Full Configuration (Recommended)
|
||||
Use this for precise control over capabilities:
|
||||
```bash
|
||||
curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- "BASE64_ENCODED_CONFIG"
|
||||
```
|
||||
|
||||
### Option B: Minimal Setup
|
||||
Use this for rapid deployment with basic telemetry:
|
||||
```bash
|
||||
curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- --url "https://sso.example.com" --token "your-host-token"
|
||||
```
|
||||
|
||||
### What this does:
|
||||
1. Downloads the latest `theta-agent` binary.
|
||||
2. Decodes the base64 configuration string into `/etc/theta42/agent.yml`.
|
||||
3. Installs a systemd service unit.
|
||||
4. Starts the agent automatically.
|
||||
|
||||
---
|
||||
|
||||
## Manual Installation
|
||||
|
||||
If you are in an air-gapped environment or prefer manual control:
|
||||
|
||||
### 1. Deploy Binary
|
||||
Place the `theta-agent` binary in `/usr/local/bin/` and ensure it is executable:
|
||||
```bash
|
||||
chmod +x /usr/local/bin/theta-agent
|
||||
```
|
||||
|
||||
### 2. Configure
|
||||
Create the configuration directory and the `agent.yml` file:
|
||||
```bash
|
||||
mkdir -p /etc/theta42
|
||||
nano /etc/theta42/agent.yml
|
||||
```
|
||||
Ensure the file has restricted permissions:
|
||||
```bash
|
||||
chmod 600 /etc/theta42/agent.yml
|
||||
```
|
||||
|
||||
### 3. Setup systemd
|
||||
Create the file `/etc/systemd/system/theta-agent.service`:
|
||||
```ini
|
||||
[Unit]
|
||||
Description=Theta Agent Unified Endpoint Management
|
||||
After=network.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStart=/usr/local/bin/theta-agent
|
||||
Restart=always
|
||||
RestartSec=5
|
||||
StandardOutput=syslog
|
||||
StandardError=syslog
|
||||
SyslogIdentifier=theta-agent
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
```
|
||||
|
||||
Enable and start the service:
|
||||
```bash
|
||||
systemctl daemon-reload
|
||||
systemctl enable theta-agent
|
||||
systemctl start theta-agent
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Verifying Connection
|
||||
Check the logs to ensure the agent has successfully connected to the SSO Manager:
|
||||
```bash
|
||||
journalctl -u theta-agent -f
|
||||
```
|
||||
You should see: `Successfully connected to SSO Manager.`
|
||||
|
||||
### Configuration Errors
|
||||
If the agent fails to start, verify the config file exists and is valid YAML:
|
||||
```bash
|
||||
ls -l /etc/theta42/agent.yml
|
||||
```
|
||||
|
||||
### Root Privileges
|
||||
The agent must run as root to execute system commands like `reboot` and `systemctl restart`. If you manually run the binary, ensure you use `sudo`.
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
# Theta Agent Protocol Specification (v1.1.0)
|
||||
|
||||
This document defines the communication protocol between the `theta-agent` (Client) and the `sso-manager` (Server).
|
||||
|
||||
## 1. Connection Establishment
|
||||
|
||||
The agent establishes a persistent outbound WebSocket connection.
|
||||
|
||||
- **Endpoint**: `wss://<manager-url>/api/agent/ws`
|
||||
- **Authentication**: The agent must provide a unique host token as a query parameter:
|
||||
- `wss://<manager-url>/api/agent/ws?token=<HOST_TOKEN>`
|
||||
|
||||
## 2. Message Format
|
||||
|
||||
All messages are exchanged as JSON objects following the `WSMessage` structure.
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "string",
|
||||
"payload": {
|
||||
"key": "value"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## 3. Client $\rightarrow$ Server Messages
|
||||
|
||||
### 3.1 Discovery (One-time & On-Change)
|
||||
Sent immediately upon connection and whenever the agent detects a change in its own network IP addresses.
|
||||
|
||||
- **Type**: `discovery`
|
||||
- **Payload**:
|
||||
- `hostname`: (string) System hostname.
|
||||
- `ip_addresses`: (array of strings) List of all non-loopback IPv4 addresses.
|
||||
- `os`: (string) OS and Platform.
|
||||
- `kernel`: (string) Kernel version.
|
||||
- `cpu`: (string) CPU model.
|
||||
- `ram_total_gb`: (float) Total system RAM in GB.
|
||||
- `disk_total_gb`: (float) Total root disk capacity in GB.
|
||||
- `location`: (string) Physical location from config.
|
||||
|
||||
### 3.2 Telemetry (Periodic)
|
||||
Sent every 30 seconds.
|
||||
|
||||
- **Type**: `telemetry`
|
||||
- **Payload**:
|
||||
- `cpu_usage_percent`: (float) Current CPU load.
|
||||
- `ram_usage_percent`: (float) Current RAM utilization.
|
||||
- `disk_usage_percent`: (float) Current root disk utilization.
|
||||
- `zfs_health`: (string) Primary ZFS pool status (e.g., "ONLINE").
|
||||
- `gpu_usage_percent`: (float) Average NVIDIA GPU utilization (-1.0 if unavailable).
|
||||
- `timestamp`: (string) RFC3339 timestamp.
|
||||
|
||||
### 3.3 Heartbeat (Periodic)
|
||||
Sent every 60 seconds to maintain the connection and signal health.
|
||||
|
||||
- **Type**: `heartbeat`
|
||||
- **Payload**:
|
||||
- `timestamp`: (string) RFC3339 timestamp.
|
||||
|
||||
### 3.4 Command Response
|
||||
Sent in response to any command received from the server.
|
||||
|
||||
- **Type**: `response` (Implicitly handled as the answer to a command)
|
||||
- **Payload**:
|
||||
- `status`: (string) Either `"ok"` or `"error"`.
|
||||
- `message`: (string) Human-readable result or error description.
|
||||
- `output`: (string, optional) Stdout/stderr for execution commands.
|
||||
|
||||
---
|
||||
|
||||
## 4. Server $\rightarrow$ Client Messages
|
||||
|
||||
### 4.1 Standard Commands
|
||||
These commands are executed if the corresponding capability is enabled in `agent.yml`.
|
||||
|
||||
| Command | Payload | Effect |
|
||||
| :--- | :--- | :--- |
|
||||
| `reload_config` | `{}` | Agent re-reads `/etc/theta42/agent.yml` from disk. |
|
||||
| `fetch_logs` | `{}` | Agent returns the last 100 lines of `journalctl -u theta-agent`. |
|
||||
|
||||
### 4.2 High-Risk Commands (Signed)
|
||||
These commands **require** an Ed25519 signature in the payload. The agent verifies the signature against the `public_key` in its config.
|
||||
|
||||
**Signature Format**:
|
||||
- The `signature` field contains the base64-encoded Ed25519 signature of the payload (with the `signature` key removed).
|
||||
|
||||
| Command | Payload | Effect |
|
||||
| :--- | :--- | :--- |
|
||||
| `reboot` | `{ "signature": "..." }` | Triggers system reboot. |
|
||||
| `service_restart` | `{ "service": "...", "signature": "..." }` | Restarts specific systemd service. |
|
||||
| `configure_ldap` | `{ "config": "...", "signature": "..." }` | Writes `/etc/sssd/sssd.conf` and restarts `sssd`. |
|
||||
| `arbitrary_bash` | `{ "script": "...", "signature": "..." }` | Executes raw bash script. |
|
||||
| `update_binary` | `{ "url": "...", "sha256": "...", "signature": "..." }` | Downloads, verifies, and replaces the agent binary. |
|
||||
|
||||
## 5. Cryptographic Verification Process
|
||||
|
||||
To send a high-risk command:
|
||||
1. Create the payload (e.g., `{"script": "uptime"}`).
|
||||
2. Canonicalize the JSON (sort keys alphabetically, remove whitespace).
|
||||
3. Sign the canonical bytes using the private Ed25519 key.
|
||||
4. Add the base64 signature to the payload: `{"script": "uptime", "signature": "..."}`.
|
||||
5. Send as a `WSMessage`.
|
||||
|
||||
The agent performs the reverse process to verify authenticity before execution.
|
||||
@@ -1,44 +1,86 @@
|
||||
# Theta Agent
|
||||
|
||||
Theta Agent is a unified endpoint management daemon for the theta42 stack. It replaces legacy bash installation scripts (like `ldap-client`) and one-way metric scripts (`telemetry-agent`) with a powerful, 2-way Command & Control (C2) Go daemon.
|
||||
Theta Agent is a unified endpoint management daemon for the theta42 stack. It replaces legacy bash installation scripts and one-way metric scripts with a powerful, 2-way Command & Control (C2) Go daemon.
|
||||
|
||||
The agent dials out to the central SSO Manager via a persistent WebSocket connection, enabling:
|
||||
- **Continuous Telemetry:** Streams CPU/RAM/ZFS/GPU health to the central inventory.
|
||||
- **Dynamic Discovery:** Automatically updates host IP and metadata on changes.
|
||||
- **Remote Operations:** Allows SSO Manager administrators to remotely configure LDAP, restart systemd services, or execute maintenance scripts.
|
||||
The agent dials out to the central SSO Manager via a persistent WebSocket connection, enabling real-time telemetry, dynamic discovery, and secure remote operations.
|
||||
|
||||
## Core Functionality
|
||||
|
||||
### 1. Telemetry & Observability
|
||||
- **Host Discovery**: Pushes a comprehensive profile (IPs, OS, Kernel, CPU, RAM/Disk) upon connection and automatically updates when network interface IPs change.
|
||||
- **Continuous Monitoring**: Streams metrics every 30 seconds:
|
||||
- CPU, RAM, and Root Disk usage.
|
||||
- **ZFS Health**: Monitors pool status via `zpool list`.
|
||||
- **GPU Utilization**: Tracks NVIDIA GPU usage via `nvidia-smi`.
|
||||
- **Health Checks**: Sends a periodic heartbeat to the SSO Manager to signal agent viability.
|
||||
|
||||
### 2. Remote Operations (C2)
|
||||
The agent provides a powerful set of administrative tools, categorized by risk:
|
||||
|
||||
#### Standard Operations
|
||||
- **Config Reload**: Triggers a reload of `/etc/theta42/agent.yml` from disk without restarting the process.
|
||||
- **Log Streaming**: Fetch the last 100 lines of the agent's system logs via the C2 channel.
|
||||
|
||||
#### High-Risk Operations (Require Cryptographic Signatures)
|
||||
To prevent unauthorized execution, these commands must be signed with a private key corresponding to the `public_key` in `agent.yml`:
|
||||
- **Service Control**: Restart approved systemd services.
|
||||
- **System Control**: Trigger a full system reboot.
|
||||
- **Config Management**: Update `/etc/sssd/sssd.conf` and restart `sssd`.
|
||||
- **Remote Execution**: Execute raw bash scripts.
|
||||
- **Self-Update**: Securely download, verify (SHA256), and apply a new binary version.
|
||||
|
||||
## The Security Model (Blast Radius & Zero-Trust)
|
||||
|
||||
Because Theta Agent runs as `root` (required to configure `/etc/sssd/sssd.conf`, restart services, and read hardware sensors), it represents a high-value target. If the central SSO Manager were compromised, a naive agent would allow an attacker to gain root shell execution on every server in the fleet.
|
||||
Because Theta Agent runs as `root`, it is a high-value target. To prevent lateral movement and contain the blast radius, it operates on a **strict, local-first capability matrix**.
|
||||
|
||||
To prevent lateral movement and contain the blast radius, **Theta Agent operates on a strict, local-first capability matrix.**
|
||||
### Local Configuration Wins
|
||||
The agent will **only** execute commands that are explicitly enabled in its local configuration file (`/etc/theta42/agent.yml`). The central SSO Manager cannot override these settings.
|
||||
|
||||
### 1. Local Configuration Wins
|
||||
The agent will **only** execute commands that are explicitly enabled in its local configuration file (`/etc/theta/agent.yml`).
|
||||
- By default, the agent is locked down to read-only telemetry and basic LDAP configuration.
|
||||
- The central SSO Manager cannot override these settings. An administrator must physically (or via local config management) edit the local `agent.yml` file to grant the agent more permissions.
|
||||
### Cryptographic Hardening
|
||||
All high-risk commands require an Ed25519 signature. The agent verifies the signature against the `public_key` provided in the local config. If the signature is missing or invalid, the command is rejected regardless of the capability matrix.
|
||||
|
||||
### 2. The Capability Matrix
|
||||
Capabilities are segmented into modules. You only enable what a specific server needs:
|
||||
### Capability Matrix
|
||||
|
||||
| Capability | Risk Level | Description |
|
||||
|------------|------------|-------------|
|
||||
| `telemetry` | Safe | Read-only. Pushes system metrics back to the SSO Manager. |
|
||||
| `configure_ldap` | Moderate | Allows the SSO manager to push down an updated SSSD configuration file. |
|
||||
| `reboot` | High | Allows the SSO Manager to trigger a system reboot. |
|
||||
| `service_control` | High | Allows starting/stopping/restarting systemd services. **Must be scoped** to specific services (e.g., `['gitea', 'nginx']`). |
|
||||
| `arbitrary_bash` | CRITICAL | Allows the execution of raw bash scripts sent from the SSO Manager. Useful for GitOps deployments on worker nodes, but highly dangerous. |
|
||||
| Capability | Risk Level | Description | Impact |
|
||||
|------------|------------|-------------|---------|
|
||||
| `telemetry` | Safe | Read-only metrics. | Pushes system health to SSO Manager. |
|
||||
| `configure_ldap` | Moderate | Configures SSSD. | Updates `/etc/sssd/sssd.conf` and restarts `sssd`. |
|
||||
| `reboot` | High | System reboot. | Triggers an immediate host reboot. |
|
||||
| `service_control` | High | Service management. | Restarts services listed in the allowed list. |
|
||||
| `arbitrary_bash` | CRITICAL | Raw bash execution. | Executes any script sent by the manager as root. |
|
||||
|
||||
### 3. Outbound-Only Communication
|
||||
The agent does not open any listening ports on the host firewall. It uses a long-lived outbound WebSocket connection to the SSO Manager.
|
||||
## Configuration
|
||||
|
||||
### 4. Cryptographic Authentication
|
||||
Every agent is issued a unique, long-lived host token during installation. The SSO Manager verifies this token to ensure commands are only routed to the intended host, and telemetry is properly attributed.
|
||||
Configuration is stored in YAML format at `/etc/theta42/agent.yml`.
|
||||
|
||||
## Example Configuration
|
||||
|
||||
See `agent.yml.example` for a secure baseline configuration.
|
||||
### Example `agent.yml`
|
||||
```yaml
|
||||
server_url: "wss://sso.theta42.local"
|
||||
auth_token: "your-unique-host-token"
|
||||
public_key: "base64-encoded-ed25519-public-key"
|
||||
location: "dc-01-rack-12"
|
||||
capabilities:
|
||||
telemetry: true
|
||||
configure_ldap: true
|
||||
reboot: true
|
||||
service_control: ["nginx", "gitea", "sssd"]
|
||||
arbitrary_bash: false
|
||||
```
|
||||
|
||||
## Installation
|
||||
|
||||
*(Coming soon: Build instructions and `theta-agent install` guide)*
|
||||
1. **Build**: Compile for your target architecture (see CI/CD artifacts).
|
||||
2. **Deploy**: Place the binary in `/usr/local/bin/theta-agent`.
|
||||
3. **Configure**: Create `/etc/theta42/agent.yml` with the required token and capabilities.
|
||||
4. **Service**: Set up as a systemd unit (example: `/etc/systemd/system/theta-agent.service`).
|
||||
|
||||
For the fastest deployment, use the installation script:
|
||||
```bash
|
||||
curl -fsSL https://sso.example.com/resources/theta-agent/install.sh | sh -s -- "BASE64_ENCODED_CONFIG"
|
||||
```
|
||||
|
||||
## Development & Testing
|
||||
|
||||
The agent uses a decoupled execution engine for safety and testability.
|
||||
- Run unit tests: `go test -v ./...`
|
||||
- The test suite uses a `MockExecutor` to verify that system commands are only triggered when the corresponding capability is enabled in the configuration.
|
||||
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
# theta-agent configuration file
|
||||
# Default location: /etc/theta/agent.yml
|
||||
# Default location: /etc/theta42/agent.yml
|
||||
|
||||
server_url: "https://sso.example.com"
|
||||
auth_token: "REPLACE_WITH_AGENT_TOKEN"
|
||||
|
||||
@@ -3,6 +3,7 @@ package main
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
@@ -19,9 +20,47 @@ type Config struct {
|
||||
ServerURL string `yaml:"server_url"`
|
||||
AuthToken string `yaml:"auth_token"`
|
||||
Location string `yaml:"location"`
|
||||
PublicKey string `yaml:"public_key"` // Ed25519 public key for signed commands
|
||||
Capabilities Capabilities `yaml:"capabilities"`
|
||||
}
|
||||
|
||||
// ConfigManager handles thread-safe access and reloading of the agent configuration.
|
||||
type ConfigManager struct {
|
||||
mu sync.RWMutex
|
||||
current *Config
|
||||
configPath string
|
||||
}
|
||||
|
||||
func NewConfigManager(path string) (*ConfigManager, error) {
|
||||
cfg, err := LoadConfig(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ConfigManager{
|
||||
current: cfg,
|
||||
configPath: path,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Get returns a copy of the current configuration.
|
||||
func (cm *ConfigManager) Get() *Config {
|
||||
cm.mu.RLock()
|
||||
defer cm.mu.RUnlock()
|
||||
return cm.current
|
||||
}
|
||||
|
||||
// Reload re-reads the configuration from disk and updates the active config.
|
||||
func (cm *ConfigManager) Reload() error {
|
||||
cfg, err := LoadConfig(cm.configPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("reload failed: %w", err)
|
||||
}
|
||||
cm.mu.Lock()
|
||||
cm.current = cfg
|
||||
cm.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func LoadConfig(path string) (*Config, error) {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoadConfig(t *testing.T) {
|
||||
// Create a temporary directory for config files
|
||||
tmpDir, err := os.MkdirTemp("", "agent-config-test")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
yamlContent string
|
||||
filename string
|
||||
expectErr bool
|
||||
}{
|
||||
{
|
||||
name: "valid config",
|
||||
yamlContent: `
|
||||
server_url: "http://sso.local"
|
||||
auth_token: "secret-token"
|
||||
location: "datacenter-1"
|
||||
capabilities:
|
||||
telemetry: true
|
||||
configure_ldap: true
|
||||
reboot: false
|
||||
service_control: ["nginx", "gitea"]
|
||||
arbitrary_bash: false
|
||||
`,
|
||||
filename: "valid.yml",
|
||||
expectErr: false,
|
||||
},
|
||||
{
|
||||
name: "invalid yaml",
|
||||
yamlContent: "invalid: [yaml: content",
|
||||
filename: "invalid.yml",
|
||||
expectErr: true,
|
||||
},
|
||||
{
|
||||
name: "missing file",
|
||||
yamlContent: "",
|
||||
filename: "nonexistent.yml",
|
||||
expectErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
path := filepath.Join(tmpDir, tc.filename)
|
||||
if tc.yamlContent != "" {
|
||||
err := os.WriteFile(path, []byte(tc.yamlContent), 0644)
|
||||
if err != nil {
|
||||
t.Fatalf("failed to write temp file: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
cfg, err := LoadConfig(path)
|
||||
if (err != nil) != tc.expectErr {
|
||||
t.Errorf("LoadConfig() error = %v, expectErr %v", err, tc.expectErr)
|
||||
return
|
||||
}
|
||||
|
||||
if !tc.expectErr && cfg == nil {
|
||||
t.Error("LoadConfig() returned nil config without error")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanManageService(t *testing.T) {
|
||||
caps := Capabilities{
|
||||
ServiceControl: []string{"nginx", "gitea"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
service string
|
||||
expected bool
|
||||
}{
|
||||
{"nginx", true},
|
||||
{"gitea", true},
|
||||
{"ssh", false},
|
||||
{"", false},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.service, func(t *testing.T) {
|
||||
if got := caps.CanManageService(tc.service); got != tc.expected {
|
||||
t.Errorf("CanManageService(%q) = %v, want %v", tc.service, got, tc.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
version: '3.8'
|
||||
services:
|
||||
theta-agent:
|
||||
build:
|
||||
context: .
|
||||
container_name: theta-agent-test
|
||||
volumes:
|
||||
- ./agent.yml.example:/etc/theta42/agent.yml:ro
|
||||
privileged: true # Required for systemctl and root operations
|
||||
network_mode: "host"
|
||||
restart: always
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
"os/exec"
|
||||
)
|
||||
|
||||
// Executor abstracts system operations to allow for easy mocking in tests.
|
||||
type Executor interface {
|
||||
// Execute runs a system command and returns combined stdout/stderr.
|
||||
Execute(command string, args ...string) ([]byte, error)
|
||||
|
||||
// WriteFile writes data to a file with specified permissions.
|
||||
WriteFile(path string, data []byte, perm os.FileMode) error
|
||||
|
||||
// ReadFile reads the content of a file.
|
||||
ReadFile(path string) ([]byte, error)
|
||||
}
|
||||
|
||||
// SystemExecutor is the production implementation that performs real system calls.
|
||||
type SystemExecutor struct{}
|
||||
|
||||
// Execute runs a real system command.
|
||||
func (s *SystemExecutor) Execute(command string, args ...string) ([]byte, error) {
|
||||
cmd := exec.Command(command, args...)
|
||||
return cmd.CombinedOutput()
|
||||
}
|
||||
|
||||
// WriteFile performs a real file write.
|
||||
func (s *SystemExecutor) WriteFile(path string, data []byte, perm os.FileMode) error {
|
||||
return os.WriteFile(path, data, perm)
|
||||
}
|
||||
|
||||
// ReadFile performs a real file read.
|
||||
func (s *SystemExecutor) ReadFile(path string) ([]byte, error) {
|
||||
return os.ReadFile(path)
|
||||
}
|
||||
@@ -3,6 +3,15 @@ module github.com/theta42/theta-agent
|
||||
go 1.22.2
|
||||
|
||||
require (
|
||||
github.com/go-ole/go-ole v1.2.6 // indirect
|
||||
github.com/gorilla/websocket v1.5.3 // indirect
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
|
||||
github.com/shirou/gopsutil/v3 v3.24.5 // indirect
|
||||
github.com/shoenig/go-m1cpu v0.1.6 // indirect
|
||||
github.com/tklauser/go-sysconf v0.3.12 // indirect
|
||||
github.com/tklauser/numcpus v0.6.1 // indirect
|
||||
github.com/yusufpapurcu/wmi v1.2.4 // indirect
|
||||
golang.org/x/sys v0.20.0 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
)
|
||||
|
||||
@@ -1,5 +1,29 @@
|
||||
github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=
|
||||
github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
|
||||
github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
|
||||
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
|
||||
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
|
||||
github.com/shirou/gopsutil/v3 v3.24.5 h1:i0t8kL+kQTvpAYToeuiVk3TgDeKOFioZO3Ztz/iZ9pI=
|
||||
github.com/shirou/gopsutil/v3 v3.24.5/go.mod h1:bsoOS1aStSs9ErQ1WWfxllSeS1K5D+U30r2NfcubMVk=
|
||||
github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM=
|
||||
github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ=
|
||||
github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
|
||||
github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
|
||||
github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=
|
||||
github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=
|
||||
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
|
||||
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
|
||||
golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
|
||||
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
|
||||
+151
@@ -0,0 +1,151 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# --- Configuration ---
|
||||
# In a real environment, these would be derived from the script's download URL
|
||||
# or passed as additional arguments. For now, we use the most recent release.
|
||||
BINARY_URL="https://github.com/theta42/theta-agent/releases/latest/download/theta-agent-linux-amd64"
|
||||
CONFIG_DIR="/etc/theta42"
|
||||
CONFIG_FILE="$CONFIG_DIR/agent.yml"
|
||||
BIN_PATH="/usr/local/bin/theta-agent"
|
||||
SERVICE_FILE="/etc/systemd/system/theta-agent.service"
|
||||
|
||||
# Colors for output
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
log() { echo -e "${GREEN}[+]${NC} $1"; }
|
||||
error() { echo -e "${RED}[!]${NC} $1"; exit 1; }
|
||||
|
||||
# 1. Root check
|
||||
if [ "$EUID" -ne 0 ]; then
|
||||
error "This script must be run as root."
|
||||
fi
|
||||
|
||||
# Install SSSD and PAM integration packages if missing
|
||||
install_sssd_deps() {
|
||||
if ! command -v sssd >/dev/null 2>&1; then
|
||||
log "Installing SSSD and PAM integration dependencies..."
|
||||
if command -v apt-get >/dev/null 2>&1; then
|
||||
DEBIAN_FRONTEND=noninteractive apt-get update -qq || true
|
||||
DEBIAN_FRONTEND=noninteractive apt-get install -y -qq sssd sssd-ldap libnss-sss libpam-sss libsss-sudo pam-auth-update || true
|
||||
if command -v pam-auth-update >/dev/null 2>&1; then
|
||||
pam-auth-update --enable mkhomedir || true
|
||||
fi
|
||||
elif command -v dnf >/dev/null 2>&1; then
|
||||
dnf install -y sssd sssd-ldap sssd-tools || true
|
||||
elif command -v yum >/dev/null 2>&1; then
|
||||
yum install -y sssd sssd-ldap sssd-tools || true
|
||||
elif command -v pacman >/dev/null 2>&1; then
|
||||
pacman -S --noconfirm sssd || true
|
||||
elif command -v zypper >/dev/null 2>&1; then
|
||||
zypper in -y sssd || true
|
||||
fi
|
||||
else
|
||||
log "SSSD is already installed."
|
||||
fi
|
||||
}
|
||||
|
||||
# 2. Argument Parsing
|
||||
URL=""
|
||||
TOKEN=""
|
||||
B64_CONFIG=""
|
||||
INSTALL_SSSD=0
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--url)
|
||||
URL="$2"
|
||||
shift 2
|
||||
;;
|
||||
--token)
|
||||
TOKEN="$2"
|
||||
shift 2
|
||||
;;
|
||||
--install-sssd|--ldap)
|
||||
INSTALL_SSSD=1
|
||||
shift
|
||||
;;
|
||||
*)
|
||||
B64_CONFIG="$1"
|
||||
shift
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Validation
|
||||
if [ -z "$B64_CONFIG" ] && [ -z "$URL" ] || [ -z "$B64_CONFIG" ] && [ -z "$TOKEN" ]; then
|
||||
error "Missing required configuration. Either provide a base64 encoded config, or both --url and --token."
|
||||
echo "Usage examples:"
|
||||
echo " sh install.sh \"BASE64_CONFIG\""
|
||||
echo " sh install.sh --url \"https://sso.local\" --token \"secret-token\" --install-sssd"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
log "Starting Theta Agent installation..."
|
||||
|
||||
# 3. Install binary
|
||||
log "Downloading binary from $BINARY_URL..."
|
||||
curl -fsSL "$BINARY_URL" -o "$BIN_PATH" || error "Failed to download binary."
|
||||
chmod +x "$BIN_PATH"
|
||||
|
||||
# 4. Setup configuration
|
||||
log "Preparing configuration directory $CONFIG_DIR..."
|
||||
mkdir -p "$CONFIG_DIR"
|
||||
chmod 755 "$CONFIG_DIR"
|
||||
|
||||
if [ -n "$B64_CONFIG" ]; then
|
||||
log "Decoding and writing configuration from base64..."
|
||||
echo "$B64_CONFIG" | base64 -d > "$CONFIG_FILE" || error "Failed to decode base64 configuration."
|
||||
else
|
||||
log "Generating minimal configuration from arguments..."
|
||||
# Create a minimal yaml with the provided URL and Token
|
||||
cat <<EOF > "$CONFIG_FILE"
|
||||
server_url: "$URL"
|
||||
auth_token: "$TOKEN"
|
||||
location: "unknown"
|
||||
capabilities:
|
||||
telemetry: true
|
||||
configure_ldap: false
|
||||
reboot: false
|
||||
service_control: []
|
||||
arbitrary_bash: false
|
||||
EOF
|
||||
fi
|
||||
chmod 600 "$CONFIG_FILE"
|
||||
|
||||
# 4b. Ensure SSSD dependencies are installed if configure_ldap is enabled
|
||||
if [ "$INSTALL_SSSD" -eq 1 ] || grep -q -i "configure_ldap:\s*true" "$CONFIG_FILE" 2>/dev/null; then
|
||||
install_sssd_deps
|
||||
fi
|
||||
|
||||
# 5. Setup systemd service
|
||||
log "Creating systemd service unit..."
|
||||
cat <<EOF > "$SERVICE_FILE"
|
||||
[Unit]
|
||||
Description=Theta Agent Unified Endpoint Management
|
||||
After=network.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
ExecStart=$BIN_PATH
|
||||
Restart=always
|
||||
RestartSec=5
|
||||
StandardOutput=syslog
|
||||
StandardError=syslog
|
||||
SyslogIdentifier=theta-agent
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
EOF
|
||||
|
||||
# 6. Start the agent
|
||||
log "Enabling and starting Theta Agent..."
|
||||
systemctl daemon-reload
|
||||
systemctl enable theta-agent
|
||||
systemctl start theta-agent
|
||||
|
||||
log "Theta Agent installation complete!"
|
||||
log "Verify status with: systemctl status theta-agent"
|
||||
log "Check logs with: journalctl -u theta-agent -f"
|
||||
@@ -12,26 +12,30 @@ func main() {
|
||||
log.Println("Starting Theta Agent...")
|
||||
|
||||
// Attempt to load configuration
|
||||
configPath := "/etc/theta/agent.yml"
|
||||
configPath := "/etc/theta42/agent.yml"
|
||||
if len(os.Args) > 1 {
|
||||
configPath = os.Args[1]
|
||||
}
|
||||
|
||||
cfg, err := LoadConfig(configPath)
|
||||
cm, err := NewConfigManager(configPath)
|
||||
if err != nil {
|
||||
log.Fatalf("Error loading configuration from %s: %v", configPath, err)
|
||||
}
|
||||
cfg := cm.Get()
|
||||
|
||||
log.Printf("Connecting to SSO Manager at %s", cfg.ServerURL)
|
||||
log.Printf("Loaded capabilities: Telemetry=%v, LDAP=%v, Reboot=%v, Bash=%v",
|
||||
cfg.Capabilities.Telemetry,
|
||||
cfg.Capabilities.ConfigureLDAP,
|
||||
cfg.Capabilities.Reboot,
|
||||
log.Printf("Loaded capabilities: Telemetry=%v, LDAP=%v, Reboot=%v, Bash=%v",
|
||||
cfg.Capabilities.Telemetry,
|
||||
cfg.Capabilities.ConfigureLDAP,
|
||||
cfg.Capabilities.Reboot,
|
||||
cfg.Capabilities.ArbitraryBash,
|
||||
)
|
||||
|
||||
// Initialize system executor
|
||||
exec := &SystemExecutor{}
|
||||
|
||||
// WebSocket connection to SSO Manager
|
||||
go connectWebSocket(cfg)
|
||||
go connectWebSocket(cm, exec)
|
||||
|
||||
// Block until signal is received
|
||||
sigs := make(chan os.Signal, 1)
|
||||
|
||||
+216
@@ -0,0 +1,216 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"net"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/shirou/gopsutil/v3/cpu"
|
||||
"github.com/shirou/gopsutil/v3/disk"
|
||||
"github.com/shirou/gopsutil/v3/host"
|
||||
"github.com/shirou/gopsutil/v3/mem"
|
||||
)
|
||||
|
||||
type DiscoveryData struct {
|
||||
Hostname string `json:"hostname"`
|
||||
IPs []string `json:"ip_addresses"`
|
||||
OS string `json:"os"`
|
||||
Kernel string `json:"kernel"`
|
||||
CPUModel string `json:"cpu"`
|
||||
RAMTotalGB float64 `json:"ram_total_gb"`
|
||||
DiskTotalGB float64 `json:"disk_total_gb"`
|
||||
Location string `json:"location"`
|
||||
}
|
||||
|
||||
type TelemetryData struct {
|
||||
CPUUsagePercent float64 `json:"cpu_usage_percent"`
|
||||
RAMUsagePercent float64 `json:"ram_usage_percent"`
|
||||
DiskUsagePercent float64 `json:"disk_usage_percent"`
|
||||
ZFSHealth string `json:"zfs_health,omitempty"`
|
||||
GPUUsage float64 `json:"gpu_usage_percent,omitempty"`
|
||||
Timestamp string `json:"timestamp"`
|
||||
}
|
||||
|
||||
// CollectDiscoveryData gathers static host information.
|
||||
func CollectDiscoveryData(cfg *Config) DiscoveryData {
|
||||
h, _ := host.Info()
|
||||
|
||||
var ips []string
|
||||
addrs, _ := net.InterfaceAddrs()
|
||||
for _, addr := range addrs {
|
||||
if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
|
||||
if ipnet.IP.To4() != nil {
|
||||
ips = append(ips, ipnet.IP.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
vm, _ := mem.VirtualMemory()
|
||||
d, _ := disk.Usage("/")
|
||||
|
||||
cpuInfo, _ := cpu.Info()
|
||||
cpuModel := "Unknown"
|
||||
if len(cpuInfo) > 0 {
|
||||
cpuModel = cpuInfo[0].Model
|
||||
}
|
||||
|
||||
return DiscoveryData{
|
||||
Hostname: h.Hostname,
|
||||
IPs: ips,
|
||||
OS: fmt.Sprintf("%s %s", h.OS, h.Platform),
|
||||
Kernel: h.KernelVersion,
|
||||
CPUModel: cpuModel,
|
||||
RAMTotalGB: float64(vm.Total) / (1024 * 1024 * 1024),
|
||||
DiskTotalGB: float64(d.Total) / (1024 * 1024 * 1024),
|
||||
Location: cfg.Location,
|
||||
}
|
||||
}
|
||||
|
||||
// CollectTelemetryData gathers real-time performance metrics including ZFS and GPU.
|
||||
func CollectTelemetryData(exec Executor) TelemetryData {
|
||||
cpuPerc, _ := cpu.Percent(time.Second, false)
|
||||
vm, _ := mem.VirtualMemory()
|
||||
d, _ := disk.Usage("/")
|
||||
|
||||
cpuVal := 0.0
|
||||
if len(cpuPerc) > 0 {
|
||||
cpuVal = cpuPerc[0]
|
||||
}
|
||||
|
||||
return TelemetryData{
|
||||
CPUUsagePercent: cpuVal,
|
||||
RAMUsagePercent: vm.UsedPercent,
|
||||
DiskUsagePercent: d.UsedPercent,
|
||||
ZFSHealth: collectZFSHealth(exec),
|
||||
GPUUsage: collectGPUUsage(exec),
|
||||
Timestamp: time.Now().Format(time.RFC3339),
|
||||
}
|
||||
}
|
||||
|
||||
func collectZFSHealth(exec Executor) string {
|
||||
out, err := exec.Execute("zpool", "list", "-H", "-o", "health")
|
||||
if err != nil {
|
||||
return "unknown"
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(string(out)), "\n")
|
||||
if len(lines) > 0 {
|
||||
return lines[0]
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
func collectGPUUsage(exec Executor) float64 {
|
||||
out, err := exec.Execute("nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits")
|
||||
if err != nil {
|
||||
return -1.0
|
||||
}
|
||||
var usage float64
|
||||
fmt.Sscanf(strings.TrimSpace(string(out)), "%f", &usage)
|
||||
return usage
|
||||
}
|
||||
|
||||
// StartTelemetryLoop manages the initial discovery push and the periodic telemetry stream.
|
||||
func StartTelemetryLoop(c MessageWriter, cm *ConfigManager, exec Executor, stopCh <-chan struct{}) {
|
||||
cfg := cm.Get()
|
||||
|
||||
// 1. Immediate Discovery Push
|
||||
pushDiscovery(c, cfg)
|
||||
|
||||
// If telemetry capability is disabled in agent.yml, return early after discovery
|
||||
if !cfg.Capabilities.Telemetry {
|
||||
log.Println("Telemetry capability is disabled in agent.yml; skipping telemetry stream.")
|
||||
return
|
||||
}
|
||||
|
||||
// 2. Periodic Telemetry Stream
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
go func() {
|
||||
defer ticker.Stop()
|
||||
var lastIPs []string
|
||||
for {
|
||||
select {
|
||||
case <-stopCh:
|
||||
return
|
||||
case <-ticker.C:
|
||||
currentCFG := cm.Get()
|
||||
if !currentCFG.Capabilities.Telemetry {
|
||||
continue
|
||||
}
|
||||
|
||||
// Network Change Detection
|
||||
currentIPs := collectIPs()
|
||||
if !equalSlices(lastIPs, currentIPs) {
|
||||
log.Println("Network change detected. Pushing discovery update...")
|
||||
pushDiscovery(c, currentCFG)
|
||||
lastIPs = currentIPs
|
||||
}
|
||||
|
||||
telemetry := CollectTelemetryData(exec)
|
||||
payload, _ := json.Marshal(WSMessage{
|
||||
Type: "telemetry",
|
||||
Payload: map[string]interface{}{
|
||||
"cpu_usage_percent": telemetry.CPUUsagePercent,
|
||||
"ram_usage_percent": telemetry.RAMUsagePercent,
|
||||
"disk_usage_percent": telemetry.DiskUsagePercent,
|
||||
"zfs_health": telemetry.ZFSHealth,
|
||||
"gpu_usage_percent": telemetry.GPUUsage,
|
||||
"timestamp": telemetry.Timestamp,
|
||||
},
|
||||
})
|
||||
if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
|
||||
log.Printf("Failed to stream telemetry: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func collectIPs() []string {
|
||||
var ips []string
|
||||
addrs, _ := net.InterfaceAddrs()
|
||||
for _, addr := range addrs {
|
||||
if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
|
||||
if ipnet.IP.To4() != nil {
|
||||
ips = append(ips, ipnet.IP.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
return ips
|
||||
}
|
||||
|
||||
func equalSlices(a, b []string) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func pushDiscovery(c MessageWriter, cfg *Config) {
|
||||
discovery := CollectDiscoveryData(cfg)
|
||||
discoveryPayload, _ := json.Marshal(discovery)
|
||||
|
||||
var discoveryMap map[string]interface{}
|
||||
json.Unmarshal(discoveryPayload, &discoveryMap)
|
||||
|
||||
msg := WSMessage{
|
||||
Type: "discovery",
|
||||
Payload: discoveryMap,
|
||||
}
|
||||
|
||||
payload, _ := json.Marshal(msg)
|
||||
if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
|
||||
log.Printf("Failed to send discovery data: %v", err)
|
||||
} else {
|
||||
log.Println("Discovery data pushed to SSO Manager.")
|
||||
}
|
||||
}
|
||||
Executable
BIN
Binary file not shown.
+295
-24
@@ -1,9 +1,17 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -15,8 +23,49 @@ type WSMessage struct {
|
||||
Payload map[string]interface{} `json:"payload"`
|
||||
}
|
||||
|
||||
func connectWebSocket(cfg *Config) {
|
||||
type MessageWriter interface {
|
||||
WriteMessage(messageType int, data []byte) error
|
||||
}
|
||||
|
||||
func verifySignature(cfg *Config, msg WSMessage) bool {
|
||||
if cfg.PublicKey == "" {
|
||||
log.Println("No public key configured; skipping signature verification")
|
||||
return true
|
||||
}
|
||||
|
||||
sigB64, ok := msg.Payload["signature"].(string)
|
||||
if !ok {
|
||||
log.Println("High-risk command missing signature")
|
||||
return false
|
||||
}
|
||||
|
||||
sig, err := base64.StdEncoding.DecodeString(sigB64)
|
||||
if err != nil {
|
||||
log.Printf("Invalid base64 signature: %v", err)
|
||||
return false
|
||||
}
|
||||
|
||||
// Create canonical payload for verification (remove signature key)
|
||||
payloadCopy := make(map[string]interface{})
|
||||
for k, v := range msg.Payload {
|
||||
if k != "signature" {
|
||||
payloadCopy[k] = v
|
||||
}
|
||||
}
|
||||
canonicalPayload, _ := json.Marshal(payloadCopy)
|
||||
|
||||
pubKeyBytes, err := base64.StdEncoding.DecodeString(cfg.PublicKey)
|
||||
if err != nil || len(pubKeyBytes) != ed25519.PublicKeySize {
|
||||
log.Printf("Invalid public key in config: %v", err)
|
||||
return false
|
||||
}
|
||||
|
||||
return ed25519.Verify(pubKeyBytes, canonicalPayload, sig)
|
||||
}
|
||||
|
||||
func connectWebSocket(cm *ConfigManager, exec Executor) {
|
||||
for {
|
||||
cfg := cm.Get()
|
||||
// Ensure protocol is ws/wss
|
||||
serverURL := strings.Replace(cfg.ServerURL, "http://", "ws://", 1)
|
||||
serverURL = strings.Replace(serverURL, "https://", "wss://", 1)
|
||||
@@ -39,12 +88,28 @@ func connectWebSocket(cfg *Config) {
|
||||
|
||||
log.Println("Successfully connected to SSO Manager.")
|
||||
|
||||
// Start telemetry if enabled
|
||||
var telemetryTicker *time.Ticker
|
||||
var telemetryDone chan bool
|
||||
if cfg.Capabilities.Telemetry {
|
||||
telemetryTicker, telemetryDone = startTelemetry(c, cfg)
|
||||
}
|
||||
stopCh := make(chan struct{})
|
||||
|
||||
// Start telemetry and discovery with stopCh lifecycle control
|
||||
StartTelemetryLoop(c, cm, exec, stopCh)
|
||||
|
||||
// Heartbeat loop
|
||||
go func() {
|
||||
ticker := time.NewTicker(60 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stopCh:
|
||||
return
|
||||
case <-ticker.C:
|
||||
hb := WSMessage{Type: "heartbeat", Payload: map[string]interface{}{"timestamp": time.Now().Format(time.RFC3339)}}
|
||||
payload, _ := json.Marshal(hb)
|
||||
if err := c.WriteMessage(websocket.TextMessage, payload); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Read loop
|
||||
for {
|
||||
@@ -60,48 +125,254 @@ func connectWebSocket(cfg *Config) {
|
||||
continue
|
||||
}
|
||||
|
||||
handleCommand(cfg, msg, c)
|
||||
handleCommand(cm, msg, c, exec)
|
||||
}
|
||||
|
||||
// Cleanup on disconnect
|
||||
close(stopCh)
|
||||
c.Close()
|
||||
if telemetryTicker != nil {
|
||||
telemetryTicker.Stop()
|
||||
telemetryDone <- true
|
||||
}
|
||||
|
||||
log.Println("WebSocket disconnected. Reconnecting in 5 seconds...")
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func handleCommand(cfg *Config, msg WSMessage, c *websocket.Conn) {
|
||||
func handleCommand(cm *ConfigManager, msg WSMessage, c MessageWriter, exec Executor) {
|
||||
cfg := cm.Get()
|
||||
log.Printf("Received command: %s", msg.Type)
|
||||
|
||||
sendResponse := func(status string, message string) {
|
||||
resp, _ := json.Marshal(map[string]string{"status": status, "message": message})
|
||||
c.WriteMessage(websocket.TextMessage, resp)
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case "config":
|
||||
log.Printf("Received config payload: %v", msg.Payload)
|
||||
case "reboot":
|
||||
if !cfg.Capabilities.Reboot {
|
||||
log.Println("Reboot rejected: capability disabled in agent.yml")
|
||||
case "reload_config":
|
||||
if err := cm.Reload(); err != nil {
|
||||
log.Printf("Reload failed: %v", err)
|
||||
sendResponse("error", "failed to reload config")
|
||||
} else {
|
||||
log.Println("Configuration reloaded successfully.")
|
||||
sendResponse("ok", "configuration reloaded")
|
||||
}
|
||||
case "fetch_logs":
|
||||
serviceName, _ := msg.Payload["service"].(string)
|
||||
if serviceName == "" {
|
||||
serviceName = "theta-agent"
|
||||
}
|
||||
|
||||
if serviceName != "theta-agent" && !cfg.Capabilities.CanManageService(serviceName) {
|
||||
log.Printf("Fetch logs rejected for '%s': not in allowed service list", serviceName)
|
||||
sendResponse("error", "service log fetch rejected")
|
||||
return
|
||||
}
|
||||
log.Println("Reboot capability enabled. (Simulation: rebooting system...)")
|
||||
|
||||
linesCount := 100
|
||||
if l, ok := msg.Payload["lines"].(float64); ok && l > 0 {
|
||||
linesCount = int(l)
|
||||
}
|
||||
|
||||
log.Printf("Fetching logs for service %s (%d lines)...", serviceName, linesCount)
|
||||
out, err := exec.Execute("journalctl", "-u", serviceName, "-n", fmt.Sprintf("%d", linesCount), "--no-pager")
|
||||
if err != nil {
|
||||
log.Printf("Log fetch failed: %v", err)
|
||||
sendResponse("error", "failed to fetch logs")
|
||||
return
|
||||
}
|
||||
resp := map[string]interface{}{
|
||||
"status": "ok",
|
||||
"service": serviceName,
|
||||
"logs": string(out),
|
||||
}
|
||||
respPayload, _ := json.Marshal(resp)
|
||||
c.WriteMessage(websocket.TextMessage, respPayload)
|
||||
return
|
||||
case "update_binary":
|
||||
if !verifySignature(cfg, msg) {
|
||||
sendResponse("error", "signature verification failed")
|
||||
return
|
||||
}
|
||||
if !cfg.Capabilities.ArbitraryBash {
|
||||
sendResponse("error", "update capability disabled")
|
||||
return
|
||||
}
|
||||
|
||||
urlStr, _ := msg.Payload["url"].(string)
|
||||
checksum, _ := msg.Payload["sha256"].(string)
|
||||
if urlStr == "" || checksum == "" {
|
||||
sendResponse("error", "missing url or sha256 checksum")
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("Updating binary from %s...", urlStr)
|
||||
if err := downloadAndUpdateBinary(urlStr, checksum); err != nil {
|
||||
log.Printf("Update failed: %v", err)
|
||||
sendResponse("error", fmt.Sprintf("update failed: %v", err))
|
||||
return
|
||||
}
|
||||
sendResponse("ok", "update applied successfully; restarting agent...")
|
||||
os.Exit(0)
|
||||
case "config":
|
||||
log.Printf("Received config payload: %v", msg.Payload)
|
||||
sendResponse("ok", "Configuration received")
|
||||
case "reboot":
|
||||
if !verifySignature(cfg, msg) {
|
||||
sendResponse("error", "signature verification failed")
|
||||
return
|
||||
}
|
||||
if !cfg.Capabilities.Reboot {
|
||||
log.Println("Reboot rejected: capability disabled in agent.yml")
|
||||
sendResponse("error", "reboot capability disabled")
|
||||
return
|
||||
}
|
||||
log.Printf("Executing reboot...")
|
||||
if _, err := exec.Execute("reboot"); err != nil {
|
||||
log.Printf("Reboot failed: %v", err)
|
||||
sendResponse("error", "reboot failed")
|
||||
return
|
||||
}
|
||||
sendResponse("ok", "system rebooting")
|
||||
case "service_restart":
|
||||
serviceName, ok := msg.Payload["service"].(string)
|
||||
if !ok || !cfg.Capabilities.CanManageService(serviceName) {
|
||||
log.Printf("Service restart rejected for '%s': not in allowed service list", serviceName)
|
||||
sendResponse("error", "service restart rejected")
|
||||
return
|
||||
}
|
||||
log.Printf("Restarting service %s...", serviceName)
|
||||
if _, err := exec.Execute("systemctl", "restart", serviceName); err != nil {
|
||||
log.Printf("Service restart failed: %v", err)
|
||||
sendResponse("error", "restart failed")
|
||||
return
|
||||
}
|
||||
sendResponse("ok", "service restarted")
|
||||
case "configure_ldap":
|
||||
if !verifySignature(cfg, msg) {
|
||||
sendResponse("error", "signature verification failed")
|
||||
return
|
||||
}
|
||||
if !cfg.Capabilities.ConfigureLDAP {
|
||||
log.Println("LDAP config rejected: capability disabled in agent.yml")
|
||||
sendResponse("error", "LDAP config disabled")
|
||||
return
|
||||
}
|
||||
|
||||
configData, ok := msg.Payload["config"].(string)
|
||||
if !ok {
|
||||
log.Println("LDAP config payload missing or not a string")
|
||||
sendResponse("error", "invalid config payload")
|
||||
return
|
||||
}
|
||||
|
||||
log.Println("Pushing updated SSSD configuration...")
|
||||
if err := exec.WriteFile("/etc/sssd/sssd.conf", []byte(configData), 0600); err != nil {
|
||||
log.Printf("Failed to write SSSD config: %v", err)
|
||||
sendResponse("error", "failed to write config")
|
||||
return
|
||||
}
|
||||
|
||||
log.Println("Restarting SSSD service...")
|
||||
if _, err := exec.Execute("systemctl", "restart", "sssd"); err != nil {
|
||||
log.Printf("SSSD restart failed: %v", err)
|
||||
sendResponse("error", "failed to restart sssd")
|
||||
return
|
||||
}
|
||||
sendResponse("ok", "LDAP configuration updated")
|
||||
case "arbitrary_bash":
|
||||
if !verifySignature(cfg, msg) {
|
||||
sendResponse("error", "signature verification failed")
|
||||
return
|
||||
}
|
||||
if !cfg.Capabilities.ArbitraryBash {
|
||||
log.Println("Bash execution rejected: capability disabled in agent.yml")
|
||||
sendResponse("error", "bash execution disabled")
|
||||
return
|
||||
}
|
||||
|
||||
script, ok := msg.Payload["script"].(string)
|
||||
if !ok {
|
||||
log.Println("Bash payload missing or not a string")
|
||||
sendResponse("error", "invalid script payload")
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("Executing remote script: %s", script)
|
||||
out, err := exec.Execute("bash", "-c", script)
|
||||
if err != nil {
|
||||
log.Printf("Script execution failed: %v", err)
|
||||
sendResponse("error", fmt.Sprintf("execution failed: %v", err))
|
||||
return
|
||||
}
|
||||
|
||||
resp := map[string]string{
|
||||
"status": "ok",
|
||||
"output": string(out),
|
||||
}
|
||||
respPayload, _ := json.Marshal(resp)
|
||||
c.WriteMessage(websocket.TextMessage, respPayload)
|
||||
return
|
||||
// heartbeat_ack is the server's acknowledgement of the agent's own periodic
|
||||
// heartbeat (the agent sends `heartbeat`, the server answers `heartbeat_ack`).
|
||||
// There is nothing to do with it -- it is not a command to run, and answering
|
||||
// an ack with an error response would inject spurious errors into the
|
||||
// command-response channel every minute. Silently ignore.
|
||||
case "heartbeat_ack":
|
||||
return
|
||||
default:
|
||||
log.Printf("Unknown command type: %s", msg.Type)
|
||||
sendResponse("error", "unknown command type")
|
||||
}
|
||||
}
|
||||
|
||||
func startTelemetry(c *websocket.Conn, cfg *Config) (*time.Ticker, chan bool) {
|
||||
ticker := time.Ticker{C: nil} // placeholder logic
|
||||
done := make(chan bool)
|
||||
log.Println("Telemetry loop started (placeholder).")
|
||||
return &ticker, done
|
||||
func downloadAndUpdateBinary(downloadURL string, expectedSHA256 string) error {
|
||||
resp, err := http.Get(downloadURL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("http fetch failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("unexpected http status: %s", resp.Status)
|
||||
}
|
||||
|
||||
tmpFile, err := os.CreateTemp("", "theta-agent-update-*")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create temp file: %w", err)
|
||||
}
|
||||
tmpPath := tmpFile.Name()
|
||||
defer os.Remove(tmpPath)
|
||||
|
||||
hasher := sha256.New()
|
||||
writer := io.MultiWriter(tmpFile, hasher)
|
||||
|
||||
if _, err := io.Copy(writer, resp.Body); err != nil {
|
||||
tmpFile.Close()
|
||||
return fmt.Errorf("failed to save binary: %w", err)
|
||||
}
|
||||
tmpFile.Close()
|
||||
|
||||
actualSHA256 := fmt.Sprintf("%x", hasher.Sum(nil))
|
||||
if !strings.EqualFold(actualSHA256, strings.TrimSpace(expectedSHA256)) {
|
||||
return fmt.Errorf("sha256 mismatch: expected %s, got %s", expectedSHA256, actualSHA256)
|
||||
}
|
||||
|
||||
if err := os.Chmod(tmpPath, 0755); err != nil {
|
||||
return fmt.Errorf("failed to set executable permissions: %w", err)
|
||||
}
|
||||
|
||||
selfPath, err := os.Executable()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to resolve current binary path: %w", err)
|
||||
}
|
||||
|
||||
resolvedPath, err := filepath.EvalSymlinks(selfPath)
|
||||
if err == nil {
|
||||
selfPath = resolvedPath
|
||||
}
|
||||
|
||||
if err := os.Rename(tmpPath, selfPath); err != nil {
|
||||
return fmt.Errorf("failed to replace binary: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type MockConn struct {
|
||||
Messages [][]byte
|
||||
}
|
||||
|
||||
func (m *MockConn) WriteMessage(messageType int, data []byte) error {
|
||||
m.Messages = append(m.Messages, data)
|
||||
return nil
|
||||
}
|
||||
|
||||
type MockExecutor struct {
|
||||
ExecutedCommands [][]string
|
||||
WrittenFiles map[string][]byte
|
||||
}
|
||||
|
||||
func (m *MockExecutor) Execute(command string, args ...string) ([]byte, error) {
|
||||
m.ExecutedCommands = append(m.ExecutedCommands, append([]string{command}, args...))
|
||||
return []byte("mock output"), nil
|
||||
}
|
||||
|
||||
func (m *MockExecutor) WriteFile(path string, data []byte, perm os.FileMode) error {
|
||||
if m.WrittenFiles == nil {
|
||||
m.WrittenFiles = make(map[string][]byte)
|
||||
}
|
||||
m.WrittenFiles[path] = data
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MockExecutor) ReadFile(path string) ([]byte, error) {
|
||||
if m.WrittenFiles != nil {
|
||||
if data, ok := m.WrittenFiles[path]; ok {
|
||||
return data, nil
|
||||
}
|
||||
}
|
||||
return []byte("mock file content"), nil
|
||||
}
|
||||
|
||||
func TestHandleCommand(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cfg *Config
|
||||
msg WSMessage
|
||||
expectedStatus string
|
||||
expectedCmd []string
|
||||
expectedFile string
|
||||
expectedFileCont string
|
||||
}{
|
||||
{
|
||||
name: "config command success",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "config",
|
||||
Payload: map[string]interface{}{"key": "value"},
|
||||
},
|
||||
expectedStatus: "ok",
|
||||
},
|
||||
{
|
||||
name: "reboot command allowed",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{Reboot: true},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "reboot",
|
||||
},
|
||||
expectedStatus: "ok",
|
||||
expectedCmd: []string{"reboot"},
|
||||
},
|
||||
{
|
||||
name: "reboot command denied",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{Reboot: false},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "reboot",
|
||||
},
|
||||
expectedStatus: "error",
|
||||
expectedCmd: nil,
|
||||
},
|
||||
{
|
||||
name: "service_restart allowed",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{
|
||||
ServiceControl: []string{"nginx"},
|
||||
},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "service_restart",
|
||||
Payload: map[string]interface{}{
|
||||
"service": "nginx",
|
||||
},
|
||||
},
|
||||
expectedStatus: "ok",
|
||||
expectedCmd: []string{"systemctl", "restart", "nginx"},
|
||||
},
|
||||
{
|
||||
name: "service_restart denied",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{
|
||||
ServiceControl: []string{"nginx"},
|
||||
},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "service_restart",
|
||||
Payload: map[string]interface{}{
|
||||
"service": "ssh",
|
||||
},
|
||||
},
|
||||
expectedStatus: "error",
|
||||
expectedCmd: nil,
|
||||
},
|
||||
{
|
||||
name: "configure_ldap allowed",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{ConfigureLDAP: true},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "configure_ldap",
|
||||
Payload: map[string]interface{}{
|
||||
"config": "domain = theta42.local\nserver = sso.local",
|
||||
},
|
||||
},
|
||||
expectedStatus: "ok",
|
||||
expectedFile: "/etc/sssd/sssd.conf",
|
||||
expectedFileCont: "domain = theta42.local\nserver = sso.local",
|
||||
expectedCmd: []string{"systemctl", "restart", "sssd"},
|
||||
},
|
||||
{
|
||||
name: "configure_ldap denied",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{ConfigureLDAP: false},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "configure_ldap",
|
||||
Payload: map[string]interface{}{
|
||||
"config": "domain = theta42.local",
|
||||
},
|
||||
},
|
||||
expectedStatus: "error",
|
||||
expectedCmd: nil,
|
||||
},
|
||||
{
|
||||
name: "arbitrary_bash allowed",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{ArbitraryBash: true},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "arbitrary_bash",
|
||||
Payload: map[string]interface{}{
|
||||
"script": "uptime",
|
||||
},
|
||||
},
|
||||
expectedStatus: "ok",
|
||||
expectedCmd: []string{"bash", "-c", "uptime"},
|
||||
},
|
||||
{
|
||||
name: "arbitrary_bash denied",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{ArbitraryBash: false},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "arbitrary_bash",
|
||||
Payload: map[string]interface{}{
|
||||
"script": "rm -rf /",
|
||||
},
|
||||
},
|
||||
expectedStatus: "error",
|
||||
expectedCmd: nil,
|
||||
},
|
||||
{
|
||||
name: "unknown command",
|
||||
cfg: &Config{
|
||||
Capabilities: Capabilities{},
|
||||
},
|
||||
msg: WSMessage{
|
||||
Type: "mystery_command",
|
||||
},
|
||||
expectedStatus: "error",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
mockConn := &MockConn{}
|
||||
mockExec := &MockExecutor{}
|
||||
cm := &ConfigManager{current: tc.cfg}
|
||||
handleCommand(cm, tc.msg, mockConn, mockExec)
|
||||
|
||||
if len(mockConn.Messages) != 1 {
|
||||
t.Fatalf("expected 1 response message, got %d", len(mockConn.Messages))
|
||||
}
|
||||
|
||||
var resp map[string]string
|
||||
if err := json.Unmarshal(mockConn.Messages[0], &resp); err != nil {
|
||||
t.Fatalf("failed to unmarshal response: %v", err)
|
||||
}
|
||||
|
||||
if resp["status"] != tc.expectedStatus {
|
||||
t.Errorf("expected status %q, got %q", tc.expectedStatus, resp["status"])
|
||||
}
|
||||
|
||||
if tc.expectedCmd != nil {
|
||||
if len(mockExec.ExecutedCommands) == 0 {
|
||||
t.Errorf("expected command to be executed, but none were")
|
||||
} else {
|
||||
cmd := mockExec.ExecutedCommands[0]
|
||||
if len(cmd) != len(tc.expectedCmd) {
|
||||
t.Errorf("expected command length %d, got %d", len(tc.expectedCmd), len(cmd))
|
||||
}
|
||||
for i := range cmd {
|
||||
if cmd[i] != tc.expectedCmd[i] {
|
||||
t.Errorf("expected arg %d = %q, got %q", i, tc.expectedCmd[i], cmd[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if len(mockExec.ExecutedCommands) > 0 {
|
||||
t.Errorf("expected no commands to be executed, but got %v", mockExec.ExecutedCommands)
|
||||
}
|
||||
|
||||
if tc.expectedFile != "" {
|
||||
content, ok := mockExec.WrittenFiles[tc.expectedFile]
|
||||
if !ok {
|
||||
t.Errorf("expected file %q to be written, but it wasn't", tc.expectedFile)
|
||||
} else if string(content) != tc.expectedFileCont {
|
||||
t.Errorf("expected file content %q, got %q", tc.expectedFileCont, string(content))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user