diff options
Diffstat (limited to 'clients')
| -rw-r--r-- | clients/umbrella-bot/bot.conf.example | 25 | ||||
| -rwxr-xr-x | clients/umbrella-bot/umbrella-bot.py | 631 | ||||
| -rw-r--r-- | clients/umbrella-bot/umbrella-bot.service | 16 | ||||
| -rw-r--r-- | clients/umbrella-cli/main.c | 410 |
4 files changed, 1082 insertions, 0 deletions
diff --git a/clients/umbrella-bot/bot.conf.example b/clients/umbrella-bot/bot.conf.example new file mode 100644 index 0000000..0cc07a1 --- /dev/null +++ b/clients/umbrella-bot/bot.conf.example @@ -0,0 +1,25 @@ +# /etc/umbrella/bot.conf +# umbrella-bot Matrix configuration +# chmod 600 this file — it contains your bot password + +# Your Synapse homeserver URL +homeserver = https://website.com + +# Bot Matrix user ID +user_id = @bot:website.com + +# Bot account password +password = yourpasswordhere + +# Room ID to listen in (not the alias, the internal ID) +# Find it in Element: Room Settings → Advanced → Internal room ID +room_id = !exampleid:website.com + +# Device ID — leave blank on first run, fill in after --setup +# device_id = UMBRELLA_BOT + +# Path to encryption key store (SQLite database) +store_path = /etc/umbrella/bot-store + +# Path to umbrella socket +socket_path = /run/umbrella/umbrella.sock diff --git a/clients/umbrella-bot/umbrella-bot.py b/clients/umbrella-bot/umbrella-bot.py new file mode 100755 index 0000000..001b94c --- /dev/null +++ b/clients/umbrella-bot/umbrella-bot.py @@ -0,0 +1,631 @@ +#!/usr/bin/env python3 +""" +umbrella-bot: Matrix bot client for the umbrella server manager. + +Connects to the umbrella Unix socket and bridges commands from a Matrix +room to umbrella units. Requires matrix-nio with encryption support. + +Dependencies: + pip install matrix-nio[e2e] aiofiles + +Setup: + 1. Create /etc/umbrella/bot.conf (see bot.conf.example) + 2. Run once with --setup to perform initial login and key verification + 3. Run normally or via systemd +""" + +import asyncio +import json +import os +import re +import signal +import socket +import struct +import sys +import logging +import argparse +from pathlib import Path +from typing import Optional + +from nio import ( + AsyncClient, + AsyncClientConfig, + LoginResponse, + MatrixRoom, + RoomMessageText, + SyncResponse, + crypto, +) +from nio.store import SqliteStore + +# ── Logging ─────────────────────────────────────────────────────────────────── + +logging.basicConfig( + level=logging.INFO, + format="[%(asctime)s] [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +log = logging.getLogger("umbrella-bot") + +# ── Config ──────────────────────────────────────────────────────────────────── + +DEFAULT_CONF = "/etc/umbrella/bot.conf" +DEFAULT_STORE = "/etc/umbrella/bot-store" +SOCK_PATH = "/run/umbrella/umbrella.sock" +PROTO_MAX = 65536 + +HELP_TEXT = """\ +Umbrella bot commands: + + !status <unit> — Show unit status + !cmd <unit> <command> — Send a console command + !restart <unit> — Restart the unit's systemd service + !update <unit> — Run the update action + !action <unit> <action> — Run a named action + !attach <unit> — Stream live log output to this room + !detach <unit> — Stop streaming log output + !units — List all units + !help — Show this message +""" + +# ── Umbrella socket client ──────────────────────────────────────────────────── + +class UmbrellaClient: + """ + Synchronous umbrella socket client. + We use a plain blocking socket here since commands are short-lived + and we don't need async for the request/response pattern. + For attach (streaming), we use a separate async wrapper. + """ + + def __init__(self, sock_path: str = SOCK_PATH): + self.sock_path = sock_path + + def _connect(self) -> socket.socket: + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.settimeout(10) + try: + s.connect(self.sock_path) + except FileNotFoundError: + raise RuntimeError(f"Umbrella socket not found: {self.sock_path} " + f"(is umbrelld running?)") + except ConnectionRefusedError: + raise RuntimeError("Umbrella daemon is not running") + return s + + def _send(self, s: socket.socket, msg: dict) -> None: + data = json.dumps(msg).encode() + header = struct.pack(">I", len(data)) + s.sendall(header + data) + + def _recv(self, s: socket.socket) -> dict: + header = b"" + while len(header) < 4: + chunk = s.recv(4 - len(header)) + if not chunk: + raise RuntimeError("Connection closed by umbrella daemon") + header += chunk + length = struct.unpack(">I", header)[0] + if length > PROTO_MAX: + raise RuntimeError(f"Message too large: {length}") + data = b"" + while len(data) < length: + chunk = s.recv(length - len(data)) + if not chunk: + raise RuntimeError("Connection closed mid-message") + data += chunk + return json.loads(data.decode()) + + def command(self, msg: dict) -> dict: + """Send a command and return the response.""" + s = self._connect() + try: + self._send(s, msg) + return self._recv(s) + finally: + s.close() + + def list_units(self) -> list: + resp = self.command({"cmd": "list"}) + return resp.get("units", []) + + def status(self, unit: str) -> dict: + return self.command({"cmd": "status", "unit": unit}) + + def send_input(self, unit: str, data: str) -> dict: + return self.command({"cmd": "input", "unit": unit, + "data": data + "\n"}) + + def run_action(self, unit: str, action: str) -> dict: + return self.command({"cmd": "action", "unit": unit, "action": action}) + + +# ── Attach session: streams output from umbrella to a Matrix room ───────────── + +class AttachSession: + """ + Maintains a persistent connection to umbrella for a specific unit, + forwarding live output to a Matrix room. + """ + + def __init__(self, unit: str, room_id: str, + matrix_client: "AsyncClient", + sock_path: str = SOCK_PATH): + self.unit = unit + self.room_id = room_id + self.matrix_client = matrix_client + self.sock_path = sock_path + self._task: Optional[asyncio.Task] = None + self._sock: Optional[socket.socket] = None + + async def start(self): + if self._task and not self._task.done(): + return # already running + self._task = asyncio.create_task(self._stream()) + + async def stop(self): + if self._task: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + if self._sock: + try: + self._sock.close() + except Exception: + pass + self._sock = None + + async def _stream(self): + loop = asyncio.get_event_loop() + + def connect_and_attach(): + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.settimeout(10) + s.connect(self.sock_path) + # Send attach command + msg = json.dumps({"cmd": "attach", "unit": self.unit}).encode() + s.sendall(struct.pack(">I", len(msg)) + msg) + # Read ok response + hdr = b"" + while len(hdr) < 4: + hdr += s.recv(4 - len(hdr)) + length = struct.unpack(">I", hdr)[0] + resp_data = b"" + while len(resp_data) < length: + resp_data += s.recv(length - len(resp_data)) + s.settimeout(None) # non-blocking from here + s.setblocking(False) + return s + + try: + self._sock = await loop.run_in_executor(None, connect_and_attach) + except Exception as e: + log.error(f"attach: failed to connect for {self.unit}: {e}") + return + + buf = b"" + log.info(f"attach: streaming {self.unit} to room {self.room_id}") + + try: + while True: + # Read available data from socket + try: + chunk = await loop.run_in_executor( + None, lambda: self._recv_one_nonblock() + ) + if chunk is None: + await asyncio.sleep(0.05) + continue + if chunk == {}: + break # disconnected + except asyncio.CancelledError: + raise + except Exception as e: + log.error(f"attach: read error for {self.unit}: {e}") + break + + if chunk.get("type") == "output": + data = chunk.get("data", "") + history = chunk.get("history", False) + if data and not history: + # Send to Matrix room, wrapped in code block + await self.matrix_client.room_send( + self.room_id, + "m.room.message", + { + "msgtype": "m.text", + "body": f"```\n{data.rstrip()}\n```", + "format": "org.matrix.custom.html", + "formatted_body": f"<pre><code>{data.rstrip()}</code></pre>", + } + ) + except asyncio.CancelledError: + pass + finally: + if self._sock: + # Send detach + try: + msg = json.dumps({"cmd": "detach"}).encode() + self._sock.setblocking(True) + self._sock.sendall(struct.pack(">I", len(msg)) + msg) + except Exception: + pass + self._sock.close() + self._sock = None + + def _recv_one_nonblock(self) -> Optional[dict]: + """Try to read one message from the non-blocking socket. Returns None if no data yet.""" + try: + hdr = self._sock.recv(4) + if not hdr: + return {} # disconnected + if len(hdr) < 4: + return None + length = struct.unpack(">I", hdr)[0] + data = b"" + self._sock.setblocking(True) + while len(data) < length: + chunk = self._sock.recv(length - len(data)) + if not chunk: + return {} + data += chunk + self._sock.setblocking(False) + return json.loads(data.decode()) + except BlockingIOError: + return None + + +# ── Power level check ───────────────────────────────────────────────────────── + +async def get_user_power_level(client: AsyncClient, + room_id: str, user_id: str) -> int: + """ + Fetch the power level for a user in a room. + Returns 0 if the user is not found or on error. + """ + try: + resp = await client.room_get_state_event( + room_id, "m.room.power_levels", "" + ) + if hasattr(resp, "content"): + users = resp.content.get("users", {}) + default = resp.content.get("users_default", 0) + return users.get(user_id, default) + except Exception as e: + log.warning(f"Could not get power level for {user_id}: {e}") + return 0 + + +# ── Command dispatch ────────────────────────────────────────────────────────── + +async def handle_command( + client: AsyncClient, + umbrella: UmbrellaClient, + attach_sessions: dict, + room: MatrixRoom, + sender: str, + text: str, + config: dict, +) -> Optional[str]: + """ + Parse and dispatch a bot command. Returns a reply string or None. + """ + parts = text.strip().split(None, 3) + if not parts: + return None + + cmd = parts[0].lower() + + # Check power level + power = await get_user_power_level(client, room.room_id, sender) + if power < 50: + return f"❌ Permission denied. Requires power level ≥ 50 (you have {power})." + + try: + if cmd == "!help": + return HELP_TEXT + + elif cmd == "!units": + units = umbrella.list_units() + if not units: + return "No units loaded." + lines = ["**Units:**"] + for u in units: + lines.append(f" `{u['name']}` — {u['display']} [{u['state']}]") + return "\n".join(lines) + + elif cmd == "!status": + if len(parts) < 2: + return "Usage: `!status <unit>`" + unit = parts[1] + resp = umbrella.status(unit) + if resp.get("type") == "error": + return f"❌ {resp['message']}" + return ( + f"**{resp.get('display', unit)}**\n" + f"State: `{resp.get('state', 'unknown')}`" + ) + + elif cmd == "!cmd": + if len(parts) < 3: + return "Usage: `!cmd <unit> <command>`" + unit = parts[1] + command = " ".join(parts[2:]) + resp = umbrella.send_input(unit, command) + if resp.get("type") == "error": + return f"❌ {resp['message']}" + return f"✓ Sent: `{command}`" + + elif cmd == "!restart": + if len(parts) < 2: + return "Usage: `!restart <unit>`" + unit = parts[1] + resp = umbrella.run_action(unit, "restart") + if resp.get("type") == "error": + return f"❌ {resp['message']}" + return f"✓ Restart dispatched for `{unit}`" + + elif cmd == "!update": + if len(parts) < 2: + return "Usage: `!update <unit>`" + unit = parts[1] + resp = umbrella.run_action(unit, "update") + if resp.get("type") == "error": + return f"❌ {resp['message']}" + return f"✓ Update dispatched for `{unit}`" + + elif cmd == "!action": + if len(parts) < 3: + return "Usage: `!action <unit> <action>`" + unit = parts[1] + action = parts[2] + resp = umbrella.run_action(unit, action) + if resp.get("type") == "error": + return f"❌ {resp['message']}" + return f"✓ Action `{action}` dispatched for `{unit}`" + + elif cmd == "!attach": + if len(parts) < 2: + return "Usage: `!attach <unit>`" + unit = parts[1] + key = (room.room_id, unit) + if key in attach_sessions: + return f"Already streaming `{unit}` to this room." + session = AttachSession(unit, room.room_id, client) + attach_sessions[key] = session + await session.start() + return f"📡 Streaming `{unit}` output to this room. Use `!detach {unit}` to stop." + + elif cmd == "!detach": + if len(parts) < 2: + return "Usage: `!detach <unit>`" + unit = parts[1] + key = (room.room_id, unit) + if key not in attach_sessions: + return f"`{unit}` is not being streamed to this room." + await attach_sessions[key].stop() + del attach_sessions[key] + return f"⏹ Stopped streaming `{unit}`." + + except RuntimeError as e: + return f"❌ Umbrella error: {e}" + except Exception as e: + log.exception(f"Command error: {e}") + return f"❌ Internal error: {e}" + + return None + + +# ── Bot core ────────────────────────────────────────────────────────────────── + +async def run_bot(config: dict): + store_path = config.get("store_path", DEFAULT_STORE) + os.makedirs(store_path, exist_ok=True) + + client_config = AsyncClientConfig( + max_limit_exceeded=0, + max_timeouts=0, + store_sync_tokens=True, + encryption_enabled=True, + ) + + client = AsyncClient( + homeserver=config["homeserver"], + user=config["user_id"], + device_id=config.get("device_id", "UMBRELLA_BOT"), + store_path=store_path, + config=client_config, + ) + + umbrella = UmbrellaClient(config.get("socket_path", SOCK_PATH)) + attach_sessions = {} + room_id = config["room_id"] + + # Login + log.info(f"Logging in as {config['user_id']}...") + resp = await client.login( + password=config["password"], + device_name="umbrella-bot", + ) + if not isinstance(resp, LoginResponse): + log.error(f"Login failed: {resp}") + return + + log.info(f"Logged in. Device ID: {client.device_id}") + + # Load encryption keys + if client.should_upload_keys: + await client.keys_upload() + await client.keys_query() + + # Message callback + async def on_message(room: MatrixRoom, event: RoomMessageText): + # Only handle messages in our configured room + if room.room_id != room_id: + return + # Ignore our own messages + if event.sender == config["user_id"]: + return + # Only handle messages that start with ! + body = event.body.strip() + if not body.startswith("!"): + return + + log.info(f"Command from {event.sender}: {body}") + + reply = await handle_command( + client, umbrella, attach_sessions, + room, event.sender, body, config + ) + if reply: + await client.room_send( + room_id, + "m.room.message", + { + "msgtype": "m.text", + "body": reply, + "format": "org.matrix.custom.html", + "formatted_body": reply.replace("\n", "<br>"), + } + ) + + client.add_event_callback(on_message, RoomMessageText) + + # Graceful shutdown + loop = asyncio.get_event_loop() + stop_event = asyncio.Event() + + def _signal_handler(): + log.info("Shutting down...") + stop_event.set() + + loop.add_signal_handler(signal.SIGTERM, _signal_handler) + loop.add_signal_handler(signal.SIGINT, _signal_handler) + + log.info(f"Syncing, watching room {room_id}") + + # Initial sync + await client.sync(timeout=10000) + await client.keys_query() + + # Main sync loop + async def sync_loop(): + while not stop_event.is_set(): + resp = await client.sync(timeout=30000, full_state=False) + if isinstance(resp, SyncResponse): + # Upload any new keys + if client.should_upload_keys: + await client.keys_upload() + + sync_task = asyncio.create_task(sync_loop()) + + await stop_event.wait() + + sync_task.cancel() + try: + await sync_task + except asyncio.CancelledError: + pass + + # Stop all attach sessions + for session in attach_sessions.values(): + await session.stop() + + await client.close() + log.info("Bot stopped.") + + +# ── Setup mode ──────────────────────────────────────────────────────────────── + +async def run_setup(config: dict): + """ + First-run setup: login, generate device keys, print device ID. + Run this once before starting the bot normally. + """ + store_path = config.get("store_path", DEFAULT_STORE) + os.makedirs(store_path, exist_ok=True) + + client = AsyncClient( + homeserver=config["homeserver"], + user=config["user_id"], + device_id=config.get("device_id", "UMBRELLA_BOT"), + store_path=store_path, + config=AsyncClientConfig(encryption_enabled=True), + ) + + print(f"Logging in as {config['user_id']}...") + resp = await client.login( + password=config["password"], + device_name="umbrella-bot", + ) + if not isinstance(resp, LoginResponse): + print(f"Login failed: {resp}") + return + + print(f"Logged in successfully.") + print(f"Device ID: {client.device_id}") + print(f"Add this to your bot.conf: device_id = {client.device_id}") + + if client.should_upload_keys: + await client.keys_upload() + print("Encryption keys uploaded.") + + print("\nSetup complete. You can now run the bot normally.") + print("Note: you may need to verify this device from another session") + print("if your room requires verified devices.") + + await client.close() + + +# ── Config loading ──────────────────────────────────────────────────────────── + +def load_config(path: str) -> dict: + config = {} + with open(path) as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + if "=" not in line: + continue + key, _, value = line.partition("=") + config[key.strip()] = value.strip() + + required = ["homeserver", "user_id", "password", "room_id"] + for key in required: + if key not in config: + raise ValueError(f"Missing required config key: {key}") + + return config + + +# ── Entry point ─────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser(description="Umbrella Matrix bot") + parser.add_argument("-c", "--config", default=DEFAULT_CONF, + help=f"Config file (default: {DEFAULT_CONF})") + parser.add_argument("--setup", action="store_true", + help="Run first-time setup and exit") + args = parser.parse_args() + + try: + config = load_config(args.config) + except FileNotFoundError: + print(f"Config file not found: {args.config}") + print(f"Create it based on bot.conf.example") + sys.exit(1) + except ValueError as e: + print(f"Config error: {e}") + sys.exit(1) + + if args.setup: + asyncio.run(run_setup(config)) + else: + asyncio.run(run_bot(config)) + + +if __name__ == "__main__": + main() diff --git a/clients/umbrella-bot/umbrella-bot.service b/clients/umbrella-bot/umbrella-bot.service new file mode 100644 index 0000000..0473ea9 --- /dev/null +++ b/clients/umbrella-bot/umbrella-bot.service @@ -0,0 +1,16 @@ +[Unit] +Description=Umbrella Matrix Bot +After=network.target umbrella.service +Requires=umbrella.service + +[Service] +Type=simple +User=root +ExecStart=/usr/bin/python3 /usr/local/bin/umbrella-bot.py -c /etc/umbrella/bot.conf +Restart=on-failure +RestartSec=10s +StandardOutput=append:/var/log/umbrella/bot.log +StandardError=append:/var/log/umbrella/bot.log + +[Install] +WantedBy=multi-user.target diff --git a/clients/umbrella-cli/main.c b/clients/umbrella-cli/main.c new file mode 100644 index 0000000..69357e4 --- /dev/null +++ b/clients/umbrella-cli/main.c @@ -0,0 +1,410 @@ +/* + * umbrella-cli: command-line client for the umbrella daemon. + * + * Usage: + * umbrella-cli list + * umbrella-cli status <unit> + * umbrella-cli attach <unit> # interactive console session + * umbrella-cli input <unit> <cmd> # send a single command + * umbrella-cli action <unit> <action> + */ + +#include <stdio.h> +#include <stdlib.h> +#include <string.h> +#include <unistd.h> +#include <errno.h> +#include <fcntl.h> +#include <termios.h> +#include <sys/socket.h> +#include <sys/un.h> +#include <sys/select.h> +#include <arpa/inet.h> + +#define SOCK_PATH "/run/umbrella/umbrella.sock" +#define PROTO_MAX 65536 +#define HDR_SIZE 4 + +/* ── Connection ──────────────────────────────────────────────────────────── */ + +static int connect_daemon(void) { + int fd = socket(AF_UNIX, SOCK_STREAM, 0); + if (fd < 0) { perror("socket"); return -1; } + + struct sockaddr_un addr; + memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + strncpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path) - 1); + + if (connect(fd, (struct sockaddr *)&addr, sizeof(addr)) != 0) { + if (errno == ENOENT) + fprintf(stderr, "umbrella: daemon not running " + "(socket not found: %s)\n", SOCK_PATH); + else + perror("connect"); + close(fd); + return -1; + } + return fd; +} + +/* ── Protocol ────────────────────────────────────────────────────────────── */ + +static int send_msg(int fd, const char *json) { + uint32_t len = (uint32_t)strlen(json); + uint32_t net = htonl(len); + if (write(fd, &net, 4) != 4) return -1; + if (write(fd, json, len) != (ssize_t)len) return -1; + return 0; +} + +static int recv_msg(int fd, char *buf, int buf_size) { + uint32_t net; + int got = 0; + while (got < 4) { + int n = read(fd, (uint8_t*)&net + got, 4 - got); + if (n <= 0) return -1; + got += n; + } + uint32_t len = ntohl(net); + if (!len || (int)len >= buf_size) return -1; + got = 0; + while (got < (int)len) { + int n = read(fd, buf + got, len - got); + if (n <= 0) return -1; + got += n; + } + buf[len] = '\0'; + return (int)len; +} + +/* ── Minimal JSON helpers ────────────────────────────────────────────────── */ + +static int jget(const char *json, const char *key, char *out, int size) { + char needle[128]; + snprintf(needle, sizeof(needle), "\"%s\":", key); + const char *p = strstr(json, needle); + if (!p) return 0; + p += strlen(needle); + while (*p == ' ') p++; + if (*p == '"') { + p++; + int i = 0; + while (*p && *p != '"' && i < size - 1) { + if (*p == '\\' && *(p+1)) { + p++; + switch (*p) { + case 'n': out[i++] = '\n'; break; + case 'r': out[i++] = '\r'; break; + case 't': out[i++] = '\t'; break; + case '"': out[i++] = '"'; break; + default: out[i++] = *p; break; + } + p++; + } else { + out[i++] = *p++; + } + } + out[i] = '\0'; + return 1; + } else { + int i = 0; + while (*p && *p != ',' && *p != '}' && i < size - 1) + out[i++] = *p++; + out[i] = '\0'; + return i > 0; + } +} + +/* ── Commands ────────────────────────────────────────────────────────────── */ + +static int cmd_list(int fd) { + if (send_msg(fd, "{\"cmd\":\"list\"}") != 0) return 1; + + char buf[PROTO_MAX]; + if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1; + + /* Simple display: find all "name"/"display"/"state" triplets */ + printf("%-24s %-32s %s\n", "NAME", "DISPLAY", "STATE"); + printf("%-24s %-32s %s\n", + "────────────────────────", + "────────────────────────────────", + "───────"); + + /* Walk the units array manually */ + const char *p = buf; + while ((p = strstr(p, "\"name\":")) != NULL) { + char name[64] = {0}, display[128] = {0}, state[32] = {0}; + + /* Extract from this position forward */ + char snippet[512]; + strncpy(snippet, p, sizeof(snippet) - 1); + + jget(snippet, "name", name, sizeof(name)); + jget(snippet, "display", display, sizeof(display)); + jget(snippet, "state", state, sizeof(state)); + + if (name[0]) + printf("%-24s %-32s %s\n", name, display, state); + + p++; /* advance past current match */ + } + return 0; +} + +static int cmd_status(int fd, const char *unit) { + char msg[256]; + snprintf(msg, sizeof(msg), + "{\"cmd\":\"status\",\"unit\":\"%s\"}", unit); + if (send_msg(fd, msg) != 0) return 1; + + char buf[PROTO_MAX]; + if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1; + + char type[32] = {0}; + jget(buf, "type", type, sizeof(type)); + if (strcmp(type, "error") == 0) { + char errmsg[256] = {0}; + jget(buf, "message", errmsg, sizeof(errmsg)); + fprintf(stderr, "Error: %s\n", errmsg); + return 1; + } + + char display[128] = {0}, state[32] = {0}, pid[16] = {0}; + jget(buf, "display", display, sizeof(display)); + jget(buf, "state", state, sizeof(state)); + jget(buf, "pid", pid, sizeof(pid)); + + printf("Unit : %s\n", unit); + printf("Display : %s\n", display); + printf("State : %s\n", state); + if (pid[0] && strcmp(pid, "0") != 0) + printf("PID : %s\n", pid); + return 0; +} + +static int cmd_input(int fd, const char *unit, const char *input) { + char msg[1024]; + snprintf(msg, sizeof(msg), + "{\"cmd\":\"input\",\"unit\":\"%s\",\"data\":\"%s\\n\"}", + unit, input); + if (send_msg(fd, msg) != 0) return 1; + + char buf[PROTO_MAX]; + if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1; + + /* Print any output that came back */ + char type[32] = {0}; + jget(buf, "type", type, sizeof(type)); + if (strcmp(type, "output") == 0) { + char data[4096] = {0}; + jget(buf, "data", data, sizeof(data)); + printf("%s\n", data); + } else if (strcmp(type, "error") == 0) { + char errmsg[256] = {0}; + jget(buf, "message", errmsg, sizeof(errmsg)); + fprintf(stderr, "Error: %s\n", errmsg); + return 1; + } + return 0; +} + +static int cmd_action(int fd, const char *unit, const char *action) { + char msg[256]; + snprintf(msg, sizeof(msg), + "{\"cmd\":\"action\",\"unit\":\"%s\",\"action\":\"%s\"}", + unit, action); + if (send_msg(fd, msg) != 0) return 1; + + char buf[PROTO_MAX]; + if (recv_msg(fd, buf, sizeof(buf)) <= 0) return 1; + + char type[32] = {0}; + jget(buf, "type", type, sizeof(type)); + if (strcmp(type, "error") == 0) { + char errmsg[256] = {0}; + jget(buf, "message", errmsg, sizeof(errmsg)); + fprintf(stderr, "Error: %s\n", errmsg); + return 1; + } + printf("Action '%s' dispatched for unit '%s'\n", action, unit); + return 0; +} + +/* + * cmd_attach: interactive console session. + * Attaches to the unit, receives live output, sends typed input. + * Ctrl+D or Ctrl+C to detach. + */ +static int cmd_attach(int fd, const char *unit) { + char msg[256]; + snprintf(msg, sizeof(msg), + "{\"cmd\":\"attach\",\"unit\":\"%s\"}", unit); + if (send_msg(fd, msg) != 0) return 1; + + /* Set terminal to raw mode so we get input character by character */ + struct termios orig, raw; + int is_tty = isatty(STDIN_FILENO); + if (is_tty) { + tcgetattr(STDIN_FILENO, &orig); + raw = orig; + raw.c_lflag &= ~(ICANON | ECHO); + raw.c_cc[VMIN] = 1; + raw.c_cc[VTIME] = 0; + tcsetattr(STDIN_FILENO, TCSANOW, &raw); + } + + printf("── Attached to %s (Ctrl+D to detach) ──\n", unit); + char drain[PROTO_MAX]; + recv_msg(fd, drain, sizeof(drain)); + fflush(stdout); + + char input_buf[1024]; + int input_pos = 0; + int ret = 0; + + while (1) { + fd_set rfds; + FD_ZERO(&rfds); + FD_SET(fd, &rfds); /* daemon messages */ + FD_SET(STDIN_FILENO, &rfds); /* keyboard input */ + int maxfd = fd > STDIN_FILENO ? fd : STDIN_FILENO; + + int n = select(maxfd + 1, &rfds, NULL, NULL, NULL); + if (n < 0) { + if (errno == EINTR) continue; + break; + } + + /* Daemon sent us something */ + if (FD_ISSET(fd, &rfds)) { + char buf[PROTO_MAX]; + int r = recv_msg(fd, buf, sizeof(buf)); + if (r <= 0) break; + + char type[32] = {0}; + jget(buf, "type", type, sizeof(type)); + + if (strcmp(type, "output") == 0) { + char data[PROTO_MAX] = {0}; + jget(buf, "data", data, sizeof(data)); + /* Print output, restore cursor line for input */ + printf("%s", data); + fflush(stdout); + } else if (strcmp(type, "error") == 0) { + char errmsg[256] = {0}; + jget(buf, "message", errmsg, sizeof(errmsg)); + fprintf(stderr, "\nError: %s\n", errmsg); + ret = 1; + break; + } + } + + /* User typed something */ + if (FD_ISSET(STDIN_FILENO, &rfds)) { + char ch; + ssize_t r = read(STDIN_FILENO, &ch, 1); + if (r <= 0) break; + + if (ch == 4) break; /* Ctrl+D — detach */ + + if (ch == '\n' || ch == '\r') { + /* Send the buffered line to the unit */ + input_buf[input_pos] = '\0'; + printf("\n"); + + char send_buf[1200]; + /* JSON-escape the input */ + char escaped[1100]; + int j = 0; + for (int i = 0; input_buf[i] && j < 1090; i++) { + if (input_buf[i] == '"') { escaped[j++] = '\\'; escaped[j++] = '"'; } + else if (input_buf[i] == '\\') { escaped[j++] = '\\'; escaped[j++] = '\\'; } + else escaped[j++] = input_buf[i]; + } + escaped[j] = '\0'; + + snprintf(send_buf, sizeof(send_buf), + "{\"cmd\":\"input\",\"unit\":\"%s\"," + "\"data\":\"%s\\n\"}", unit, escaped); + send_msg(fd, send_buf); + input_pos = 0; + } else if (ch == 127 || ch == '\b') { + /* Backspace */ + if (input_pos > 0) { + input_pos--; + printf("\b \b"); + fflush(stdout); + } + } else { + if (input_pos < (int)sizeof(input_buf) - 1) { + input_buf[input_pos++] = ch; + /* Echo the character */ + write(STDOUT_FILENO, &ch, 1); + } + } + } + } + + /* Restore terminal */ + if (is_tty) + tcsetattr(STDIN_FILENO, TCSANOW, &orig); + + /* Send detach */ + send_msg(fd, "{\"cmd\":\"detach\"}"); + printf("\n── Detached ──\n"); + return ret; +} + +/* ── Entry point ─────────────────────────────────────────────────────────── */ + +static void usage(void) { + fprintf(stderr, + "Usage: umbrella-cli <command> [args]\n" + "\n" + "Commands:\n" + " list List all units\n" + " status <unit> Show unit status\n" + " attach <unit> Interactive console session\n" + " input <unit> <cmd> Send a single command\n" + " action <unit> <action> Run a named action\n" + "\n" + "Examples:\n" + " umbrella-cli list\n" + " umbrella-cli attach tf2-novemen\n" + " umbrella-cli input tf2-novemen \"say Server restarting soon\"\n" + " umbrella-cli action tf2-novemen update\n"); +} + +int main(int argc, char *argv[]) { + if (argc < 2) { usage(); return 1; } + + int fd = connect_daemon(); + if (fd < 0) return 1; + + int ret = 0; + const char *cmd = argv[1]; + + if (strcmp(cmd, "list") == 0) { + ret = cmd_list(fd); + } else if (strcmp(cmd, "status") == 0) { + if (argc < 3) { fprintf(stderr, "status: need unit name\n"); ret = 1; } + else ret = cmd_status(fd, argv[2]); + } else if (strcmp(cmd, "attach") == 0) { + if (argc < 3) { fprintf(stderr, "attach: need unit name\n"); ret = 1; } + else ret = cmd_attach(fd, argv[2]); + } else if (strcmp(cmd, "input") == 0) { + if (argc < 4) { fprintf(stderr, "input: need unit and command\n"); ret = 1; } + else ret = cmd_input(fd, argv[2], argv[3]); + } else if (strcmp(cmd, "action") == 0) { + if (argc < 4) { fprintf(stderr, "action: need unit and action name\n"); ret = 1; } + else ret = cmd_action(fd, argv[2], argv[3]); + } else { + fprintf(stderr, "Unknown command: %s\n\n", cmd); + usage(); + ret = 1; + } + + close(fd); + return ret; +} |