context: add context

pkh context allows to manage contexts (local, ssh) and run contextualized commands (deb)

in other words, this allows building binary packages over ssh

Cargo.toml

@@ -24,6 +24,9 @@ env_logger = "0.11.8"
 futures-util = { version = "0.3.31", features = ["tokio-io"] }
 tar = "0.4"
 xz2 = "0.1"
+serde_json = "1.0.145"
+directories = "6.0.0"
+ssh2 = "0.9.5"
 
 [dev-dependencies]
 tempfile = "3.10.1"

src/context/api.rs

@@ -0,0 +1,113 @@
+use serde::{Deserialize, Serialize};
+use std::ffi::OsStr;
+use std::io;
+use std::path::{Path, PathBuf};
+
+use super::local::LocalDriver;
+use super::ssh::SshDriver;
+
+/// Internal trait defining the strategy for executing a command.
+///
+/// This allows to have different execution behaviors (Local, SSH, ...) while keeping the
+/// `ContextCommand` API consistent.
+pub trait CommandRunner {
+    fn add_arg(&mut self, arg: String);
+    fn status(&mut self) -> io::Result<std::process::ExitStatus>;
+    fn output(&mut self) -> io::Result<std::process::Output>;
+}
+
+pub trait ContextDriver {
+    fn ensure_available(&self, src: &Path, dest_root: &str) -> io::Result<PathBuf>;
+    fn create_runner(&self, program: String) -> Box<dyn CommandRunner>;
+    fn prepare_work_dir(&self) -> io::Result<String>;
+}
+
+/// Represents an execution environment (Local or via SSH).
+///
+/// This is the data model that configuration files store. It defines *where* operations happen.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(tag = "type")]
+#[derive(Default)]
+pub enum Context {
+    #[serde(rename = "local")]
+    #[default]
+    Local,
+    #[serde(rename = "ssh")]
+    Ssh {
+        host: String,
+        user: Option<String>,
+        port: Option<u16>,
+    },
+}
+
+impl Context {
+    pub fn command<S: AsRef<OsStr>>(&self, program: S) -> ContextCommand {
+        let runner = self
+            .driver()
+            .create_runner(program.as_ref().to_string_lossy().to_string());
+        ContextCommand { runner }
+    }
+
+    /// Ensures that the source file/directory exists at the destination context.
+    ///
+    /// If Local: Returns the absolute path of `src`.
+    /// If Remote: Copies `src` to `dest_root` on the remote and returns the path to the copied entity.
+    pub fn ensure_available(&self, src: &Path, dest_root: &str) -> io::Result<PathBuf> {
+        self.driver().ensure_available(src, dest_root)
+    }
+
+    pub fn prepare_work_dir(&self) -> io::Result<String> {
+        self.driver().prepare_work_dir()
+    }
+
+    fn driver(&self) -> Box<dyn ContextDriver> {
+        match self {
+            Context::Local => Box::new(LocalDriver),
+            Context::Ssh { host, user, port } => Box::new(SshDriver {
+                host: host.clone(),
+                user: user.clone(),
+                port: *port,
+            }),
+        }
+    }
+}
+
+/// A builder and executor for commands within a specific `Context`.
+///
+/// This struct acts as a high-level facade (API) for command creation. It hides the implementation
+/// details of *how* the command is actually executed, allowing the user to simple chain arguments
+/// and call `status()` or `output()`.
+///
+/// It delegates the actual work to a `CommandRunner`.
+pub struct ContextCommand {
+    runner: Box<dyn CommandRunner>,
+}
+
+impl ContextCommand {
+    pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self {
+        self.runner
+            .add_arg(arg.as_ref().to_string_lossy().to_string());
+        self
+    }
+
+    // Support chaining args
+    pub fn args<I, S>(&mut self, args: I) -> &mut Self
+    where
+        I: IntoIterator<Item = S>,
+        S: AsRef<OsStr>,
+    {
+        for arg in args {
+            self.arg(arg);
+        }
+        self
+    }
+
+    pub fn status(&mut self) -> io::Result<std::process::ExitStatus> {
+        self.runner.status()
+    }
+
+    // Capture output
+    pub fn output(&mut self) -> io::Result<std::process::Output> {
+        self.runner.output()
+    }
+}

src/context/local.rs

@@ -0,0 +1,45 @@
+use super::api::{CommandRunner, ContextDriver};
+/// Local context: execute commands locally
+/// Context driver: Does nothing
+/// Command runner: Wrapper around 'std::process::Command'
+use std::io;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+
+pub struct LocalDriver;
+
+impl ContextDriver for LocalDriver {
+    fn ensure_available(&self, src: &Path, _dest_root: &str) -> io::Result<PathBuf> {
+        src.canonicalize()
+    }
+
+    fn create_runner(&self, program: String) -> Box<dyn CommandRunner> {
+        Box::new(LocalRunner {
+            program,
+            args: Vec::new(),
+        })
+    }
+
+    fn prepare_work_dir(&self) -> io::Result<String> {
+        Ok(".".to_string())
+    }
+}
+
+pub struct LocalRunner {
+    pub program: String,
+    pub args: Vec<String>,
+}
+
+impl CommandRunner for LocalRunner {
+    fn add_arg(&mut self, arg: String) {
+        self.args.push(arg);
+    }
+
+    fn status(&mut self) -> io::Result<std::process::ExitStatus> {
+        Command::new(&self.program).args(&self.args).status()
+    }
+
+    fn output(&mut self) -> io::Result<std::process::Output> {
+        Command::new(&self.program).args(&self.args).output()
+    }
+}

src/context/manager.rs

@@ -0,0 +1,117 @@
+use directories::ProjectDirs;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::fs;
+use std::io;
+use std::path::PathBuf;
+
+use super::api::Context;
+
+#[derive(Debug, Serialize, Deserialize, Default, Clone)]
+pub struct Config {
+    pub current_context: Option<String>,
+    pub contexts: HashMap<String, Context>,
+}
+
+pub struct ContextManager {
+    config_path: PathBuf,
+    config: Config,
+}
+
+impl ContextManager {
+    pub fn new() -> io::Result<Self> {
+        let proj_dirs = ProjectDirs::from("com", "pkh", "pkh").ok_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::NotFound,
+                "Could not determine config directory",
+            )
+        })?;
+        let config_dir = proj_dirs.config_dir();
+        fs::create_dir_all(config_dir)?;
+        let config_path = config_dir.join("contexts.json");
+
+        let config = if config_path.exists() {
+            let content = fs::read_to_string(&config_path)?;
+            serde_json::from_str(&content)
+                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
+        } else {
+            let mut cfg = Config::default();
+            cfg.contexts.insert("local".to_string(), Context::Local);
+            cfg.current_context = Some("local".to_string());
+            cfg
+        };
+
+        Ok(Self {
+            config_path,
+            config,
+        })
+    }
+
+    pub fn with_path(path: PathBuf) -> Self {
+        Self {
+            config_path: path,
+            config: Config::default(),
+        }
+    }
+
+    pub fn save(&self) -> io::Result<()> {
+        let content = serde_json::to_string_pretty(&self.config)
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+        fs::write(&self.config_path, content)?;
+        Ok(())
+    }
+
+    pub fn list_contexts(&self) -> Vec<String> {
+        self.config.contexts.keys().cloned().collect()
+    }
+
+    pub fn get_context(&self, name: &str) -> Option<&Context> {
+        self.config.contexts.get(name)
+    }
+
+    pub fn add_context(&mut self, name: &str, context: Context) -> io::Result<()> {
+        self.config.contexts.insert(name.to_string(), context);
+        self.save()
+    }
+
+    pub fn remove_context(&mut self, name: &str) -> io::Result<()> {
+        if self.config.contexts.remove(name).is_some() {
+            if self.config.current_context.as_deref() == Some(name) {
+                self.config.current_context = Some("local".to_string());
+                if !self.config.contexts.contains_key("local") {
+                    self.config
+                        .contexts
+                        .insert("local".to_string(), Context::Local);
+                }
+            }
+            self.save()?;
+        }
+        Ok(())
+    }
+
+    pub fn set_current(&mut self, name: &str) -> io::Result<()> {
+        if self.config.contexts.contains_key(name) {
+            self.config.current_context = Some(name.to_string());
+            self.save()?;
+            Ok(())
+        } else {
+            Err(io::Error::new(
+                io::ErrorKind::NotFound,
+                format!("Context '{}' not found", name),
+            ))
+        }
+    }
+
+    pub fn current(&self) -> Context {
+        self.config
+            .current_context
+            .as_deref()
+            .and_then(|name| self.config.contexts.get(name))
+            .cloned()
+            .unwrap_or(Context::Local)
+    }
+
+    pub fn current_name(&self) -> Option<String> {
+        self.config.current_context.clone()
+    }
+}

src/context/mod.rs

@@ -0,0 +1,88 @@
+mod api;
+mod local;
+mod manager;
+mod ssh;
+
+pub use api::{CommandRunner, Context, ContextCommand};
+pub use manager::ContextManager;
+
+pub fn current_context() -> Context {
+    match ContextManager::new() {
+        Ok(mgr) => mgr.current(),
+        Err(_) => Context::Local,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::fs;
+    use tempfile::NamedTempFile;
+
+    #[test]
+    fn test_ensure_available_local() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let src_file = temp_dir.path().join("src.txt");
+        fs::write(&src_file, "local").unwrap();
+
+        let ctx = Context::Local;
+        let dest = ctx.ensure_available(&src_file, "/tmp").unwrap();
+
+        // Should return canonical path
+        assert_eq!(dest, src_file.canonicalize().unwrap());
+    }
+
+    #[test]
+    fn test_context_manager_crud() {
+        let temp_file = NamedTempFile::new().unwrap();
+        let path = temp_file.path().to_path_buf();
+
+        let mut mgr = ContextManager::with_path(path.clone());
+
+        // Add
+        let ssh_ctx = Context::Ssh {
+            host: "10.0.0.1".into(),
+            user: Some("admin".into()),
+            port: Some(2222),
+        };
+        mgr.add_context("myserver", ssh_ctx.clone()).unwrap();
+
+        assert!(mgr.get_context("myserver").is_some());
+        assert_eq!(mgr.get_context("myserver").unwrap(), &ssh_ctx);
+
+        // List
+        let list = mgr.list_contexts();
+        assert!(list.contains(&"myserver".to_string()));
+
+        // Set Current
+        mgr.set_current("myserver").unwrap();
+        assert_eq!(mgr.current(), ssh_ctx);
+        assert_eq!(mgr.current_name(), Some("myserver".to_string()));
+
+        // Remove
+        mgr.remove_context("myserver").unwrap();
+        assert!(mgr.get_context("myserver").is_none());
+        assert_eq!(mgr.current(), Context::Local);
+    }
+
+    #[test]
+    fn test_persistence() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let config_path = temp_dir.path().join("contexts.json");
+
+        {
+            let mut mgr = ContextManager::with_path(config_path.clone());
+            mgr.add_context("persistent", Context::Local).unwrap();
+            mgr.set_current("persistent").unwrap();
+        }
+
+        let content = fs::read_to_string(&config_path).unwrap();
+        let loaded_config: super::manager::Config = serde_json::from_str(&content).unwrap();
+
+        assert_eq!(
+            loaded_config.current_context,
+            Some("persistent".to_string())
+        );
+        assert!(loaded_config.contexts.contains_key("persistent"));
+    }
+}

src/context/ssh.rs

@@ -0,0 +1,199 @@
+use super::api::{CommandRunner, ContextDriver};
+use log::debug;
+use std::fs;
+/// SSH context: execute commands over an SSH connection
+/// Context driver: Copies over SFTP with ssh2
+/// Command runner: Executes commands over an ssh2 channel (with PTY)
+use std::io::{self, Read};
+use std::net::TcpStream;
+#[cfg(unix)]
+use std::os::unix::process::ExitStatusExt;
+use std::path::{Path, PathBuf};
+use std::process::ExitStatus;
+
+pub fn connect_ssh(host: &str, user: Option<&str>, port: Option<u16>) -> io::Result<ssh2::Session> {
+    let port = port.unwrap_or(22);
+    let tcp = TcpStream::connect((host, port))?;
+    let mut session = ssh2::Session::new().unwrap();
+    session.set_tcp_stream(tcp);
+    session
+        .handshake()
+        .map_err(|e| io::Error::new(io::ErrorKind::ConnectionRefused, e))?;
+
+    // Username: if set, use it; else, use local 'USER', and if unset, use 'root'
+    let local_user_string;
+    let user = match user {
+        Some(u) => u,
+        None => {
+            local_user_string = std::env::var("USER").unwrap_or_else(|_| "root".to_string());
+            &local_user_string
+        }
+    };
+
+    if session.userauth_agent(user).is_ok() {
+        return Ok(session);
+    }
+
+    if !session.authenticated() {
+        return Err(io::Error::new(
+            io::ErrorKind::PermissionDenied,
+            "SSH authentication failed (tried agent)",
+        ));
+    }
+
+    Ok(session)
+}
+
+pub struct SshDriver {
+    pub host: String,
+    pub user: Option<String>,
+    pub port: Option<u16>,
+}
+
+impl ContextDriver for SshDriver {
+    fn ensure_available(&self, src: &Path, dest_root: &str) -> io::Result<PathBuf> {
+        let sess = connect_ssh(&self.host, self.user.as_deref(), self.port)?;
+        let sftp = sess.sftp().map_err(io::Error::other)?;
+
+        let file_name = src
+            .file_name()
+            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "src has no filename"))?;
+        let remote_dest = Path::new(dest_root).join(file_name);
+
+        debug!("Uploading {:?} to remote {:?}", src, remote_dest);
+
+        Self::upload_recursive(&sftp, src, &remote_dest)?;
+
+        Ok(remote_dest)
+    }
+
+    fn create_runner(&self, program: String) -> Box<dyn CommandRunner> {
+        Box::new(SshRunner {
+            host: self.host.clone(),
+            user: self.user.clone(),
+            port: self.port,
+            program,
+            args: Vec::new(),
+        })
+    }
+
+    fn prepare_work_dir(&self) -> io::Result<String> {
+        let sess = connect_ssh(&self.host, self.user.as_deref(), self.port)?;
+        let mut channel = sess.channel_session().map_err(io::Error::other)?;
+
+        channel.exec("mktemp -d").map_err(io::Error::other)?;
+
+        let mut stdout = String::new();
+        channel.read_to_string(&mut stdout)?;
+        channel.wait_close().map_err(io::Error::other)?;
+
+        if channel.exit_status().unwrap_or(-1) != 0 {
+            return Err(io::Error::other(
+                "Failed to create remote temporary directory",
+            ));
+        }
+
+        Ok(stdout.trim().to_string())
+    }
+}
+
+impl SshDriver {
+    fn upload_recursive(sftp: &ssh2::Sftp, src: &Path, dest: &Path) -> io::Result<()> {
+        if src.is_dir() {
+            // Create dir
+            let _ = sftp.mkdir(dest, 0o755);
+            for entry in fs::read_dir(src)? {
+                let entry = entry?;
+                let path = entry.path();
+                let name = entry.file_name();
+                let dest_path = dest.join(name);
+                Self::upload_recursive(sftp, &path, &dest_path)?;
+            }
+        } else {
+            let mut file = fs::File::open(src)?;
+            let mut remote_file = sftp.create(dest).map_err(|e| {
+                io::Error::other(format!("Failed to create remote file {:?}: {}", dest, e))
+            })?;
+            io::copy(&mut file, &mut remote_file)?;
+        }
+        Ok(())
+    }
+}
+
+pub struct SshRunner {
+    pub host: String,
+    pub user: Option<String>,
+    pub port: Option<u16>,
+    pub program: String,
+    pub args: Vec<String>,
+}
+
+impl SshRunner {
+    fn prepare_channel(&self) -> io::Result<(ssh2::Channel, String)> {
+        let sess = connect_ssh(&self.host, self.user.as_deref(), self.port)?;
+        let channel = sess.channel_session().map_err(io::Error::other)?;
+
+        // Construct command line
+        let mut cmd_line = self.program.clone();
+        for arg in &self.args {
+            cmd_line.push(' ');
+            cmd_line.push_str(arg); // TODO: escape
+        }
+        Ok((channel, cmd_line))
+    }
+}
+
+impl CommandRunner for SshRunner {
+    fn add_arg(&mut self, arg: String) {
+        self.args.push(arg);
+    }
+
+    fn status(&mut self) -> io::Result<ExitStatus> {
+        let (mut channel, cmd_line) = self.prepare_channel()?;
+
+        debug!("Executing SSH command: {}", cmd_line);
+
+        channel
+            .request_pty("xterm", None, None)
+            .map_err(|e| io::Error::other(format!("Failed to request PTY: {}", e)))?;
+
+        channel.exec(&cmd_line).map_err(io::Error::other)?;
+
+        let mut stdout_stream = channel.stream(0);
+        let mut stdout = io::stdout();
+        io::copy(&mut stdout_stream, &mut stdout)?;
+
+        channel.wait_close().map_err(io::Error::other)?;
+
+        let code = channel.exit_status().unwrap_or(-1);
+        Ok(ExitStatus::from_raw(code))
+    }
+
+    fn output(&mut self) -> io::Result<std::process::Output> {
+        let (mut channel, cmd_line) = self.prepare_channel()?;
+
+        channel.exec(&cmd_line).map_err(io::Error::other)?;
+
+        let mut stdout = Vec::new();
+        channel.read_to_end(&mut stdout)?;
+
+        let mut stderr = Vec::new();
+        channel.stderr().read_to_end(&mut stderr)?;
+
+        channel.wait_close().map_err(io::Error::other)?;
+
+        #[cfg(unix)]
+        let status = std::process::ExitStatus::from_raw(channel.exit_status().unwrap_or(-1));
+
+        #[cfg(not(unix))]
+        let status = {
+            panic!("SSH output capture only supported on Unix");
+        };
+
+        Ok(std::process::Output {
+            status,
+            stdout,
+            stderr,
+        })
+    }
+}

src/deb.rs

@@ -1,6 +1,6 @@
+use crate::context::{Context, current_context};
 use std::error::Error;
-use std::path::Path;
-use std::process::Command;
+use std::path::{Path, PathBuf};
 
 pub fn build_binary_package(
     arch: Option<&str>,
@@ -13,29 +13,122 @@ pub fn build_binary_package(
     let changelog_path = cwd.join("debian/changelog");
     let (package, version, _series) = crate::changelog::parse_changelog_header(&changelog_path)?;
 
-    // Construct dsc file name
+    // Find .dsc file
+    let dsc_path = find_dsc_file(cwd, &package, &version)?;
+    println!("Building {} using sbuild...", dsc_path.display());
+
+    // Identify all related files from .dsc
+    let mut files_to_ensure = get_dsc_related_files(&dsc_path)?;
+    // Ensure dsc itself is included (usually first)
+    if !files_to_ensure.contains(&dsc_path) {
+        files_to_ensure.insert(0, dsc_path.clone());
+    }
+
+    // Prepare Environment
+    let ctx = current_context();
+    let build_root = ctx.prepare_work_dir()?;
+
+    // Ensure availability of all needed files for the build
+    let remote_dsc_path = upload_package_files(&ctx, &files_to_ensure, &build_root, &dsc_path)?;
+    println!(
+        "Building {} on {}...",
+        remote_dsc_path.display(),
+        build_root
+    );
+
+    // Run sbuild
+    run_sbuild(&ctx, &remote_dsc_path, arch, series)?;
+
+    Ok(())
+}
+
+fn find_dsc_file(cwd: &Path, package: &str, version: &str) -> Result<PathBuf, Box<dyn Error>> {
     let parent = cwd.parent().ok_or("Cannot find parent directory")?;
     let dsc_name = format!("{}_{}.dsc", package, version);
     let dsc_path = parent.join(&dsc_name);
+
     if !dsc_path.exists() {
         return Err(format!("Could not find .dsc file at {}", dsc_path.display()).into());
     }
+    Ok(dsc_path)
+}
 
-    println!("Building {} using sbuild...", dsc_path.display());
+fn get_dsc_related_files(dsc_path: &Path) -> Result<Vec<PathBuf>, Box<dyn Error>> {
+    let content = std::fs::read_to_string(dsc_path)?;
+    let parent = dsc_path.parent().unwrap(); // dsc_path exists so parent exists
+    let mut files = Vec::new();
 
-    let mut status = Command::new("sbuild");
-    if let Some(arch) = arch {
-        status.arg(format!("--arch={}", arch));
+    let mut in_files = false;
+    for line in content.lines() {
+        if line.starts_with("Files:") {
+            in_files = true;
+            continue;
+        }
+        if in_files {
+            if line.starts_with(' ') {
+                let parts: Vec<&str> = line.split_whitespace().collect();
+                if parts.len() >= 3 {
+                    let filename = parts[2];
+                    let filepath = parent.join(filename);
+                    if filepath.exists() {
+                        files.push(filepath);
+                    } else {
+                        return Err(
+                            format!("Referenced file {} not found", filepath.display()).into()
+                        );
+                    }
+                }
+            } else {
+                in_files = false;
+            }
+        }
     }
-    if let Some(series) = series {
-        status.arg(format!("--dist={}", series));
+    Ok(files)
+}
+
+fn upload_package_files(
+    ctx: &Context,
+    files: &[PathBuf],
+    dest_root: &str,
+    local_dsc_path: &Path,
+) -> Result<PathBuf, Box<dyn Error>> {
+    let mut remote_dsc_path = PathBuf::new();
+
+    for file in files {
+        let remote_path = ctx.ensure_available(file, dest_root)?;
+        // Check if this is the dsc file by comparing file names
+        if let (Some(f_name), Some(dsc_name)) = (file.file_name(), local_dsc_path.file_name())
+            && f_name == dsc_name
+        {
+            remote_dsc_path = remote_path;
+        }
     }
 
-    let status = status.arg(dsc_path).status()?;
+    if remote_dsc_path.as_os_str().is_empty() {
+        return Err("Failed to determine remote path for .dsc file".into());
+    }
+
+    Ok(remote_dsc_path)
+}
+
+fn run_sbuild(
+    ctx: &Context,
+    dsc_path: &Path,
+    arch: Option<&str>,
+    series: Option<&str>,
+) -> Result<(), Box<dyn Error>> {
+    let mut cmd = ctx.command("sbuild");
+    if let Some(a) = arch {
+        cmd.arg(format!("--arch={}", a));
+    }
+    if let Some(s) = series {
+        cmd.arg(format!("--dist={}", s));
+    }
+
+    let status = cmd.arg(dsc_path).status()?;
 
     if !status.success() {
         return Err(format!("sbuild failed with status: {}", status).into());
     }
-
     Ok(())
 }

src/lib.rs

@@ -1,5 +1,6 @@
 pub mod build;
 pub mod changelog;
+pub mod context;
 pub mod deb;
 pub mod package_info;
 pub mod pull;

src/main.rs

@@ -53,7 +53,31 @@ fn main() {
             Command::new("deb")
                 .about("Build the binary package")
                 .arg(arg!(-s --series <series> "Target distribution series").required(false))
-                .arg(arg!(-a --arch <arch> "Target architecture").required(false))
+                .arg(arg!(-a --arch <arch> "Target architecture").required(false)),
+        )
+        .subcommand(
+            Command::new("context")
+                .about("Manage contexts")
+                .subcommand_required(true)
+                .subcommand(
+                    Command::new("create")
+                        .about("Create a new context")
+                        .arg(arg!(<name> "Context name"))
+                        .arg(arg!(--type <type> "Context type: ssh (only type supported for now)"))
+                        .arg(arg!(--endpoint <endpoint> "Context endpoint (for example: ssh://user@host:port)"))
+                )
+                .subcommand(
+                    Command::new("rm")
+                        .about("Remove a context")
+                        .arg(arg!(<name> "Context name"))
+                )
+                .subcommand(Command::new("ls").about("List contexts"))
+                .subcommand(Command::new("show").about("Show current context"))
+                .subcommand(
+                    Command::new("use")
+                        .about("Set current context")
+                        .arg(arg!(<name> "Context name"))
+                )
         )
         .get_matches();
 
@@ -124,6 +148,100 @@ fn main() {
                 std::process::exit(1);
             }
         }
+        Some(("context", sub_matches)) => {
+            use pkh::context::{Context, ContextManager};
+
+            let mut mgr = match ContextManager::new() {
+                Ok(mgr) => mgr,
+                Err(e) => {
+                    error!("Failed to initialize context manager: {}", e);
+                    std::process::exit(1);
+                }
+            };
+
+            match sub_matches.subcommand() {
+                Some(("create", args)) => {
+                    let name = args.get_one::<String>("name").unwrap();
+                    let type_str = args
+                        .get_one::<String>("type")
+                        .map(|s| s.as_str())
+                        .unwrap_or("local");
+
+                    let context = match type_str {
+                        "local" => Context::Local,
+                        "ssh" => {
+                            let endpoint = args
+                                .get_one::<String>("endpoint")
+                                .expect("Endpoint is required for ssh context");
+
+                            // Parse host, user, port from endpoint
+                            // Formats: [ssh://][user@]host[:port]
+                            let endpoint_re = regex::Regex::new(r"^(?:ssh://)?(?:(?P<user>[^@]+)@)?(?P<host>[^:/]+)(?::(?P<port>\d+))?$").unwrap();
+                            let endpoint_cap = endpoint_re.captures(endpoint).unwrap_or_else(|| {
+                                error!("Invalid endpoint format: '{}'. Expected [ssh://][user@]host[:port]", endpoint);
+                                std::process::exit(1);
+                            });
+
+                            let host = endpoint_cap.name("host").unwrap().as_str().to_string();
+                            let user = endpoint_cap.name("user").map(|m| m.as_str().to_string());
+                            let port = endpoint_cap.name("port").map(|m| {
+                                m.as_str().parse::<u16>().unwrap_or_else(|_| {
+                                    error!("Invalid port number");
+                                    std::process::exit(1);
+                                })
+                            });
+
+                            Context::Ssh { host, user, port }
+                        }
+                        _ => {
+                            error!("Unknown context type: {}", type_str);
+                            std::process::exit(1);
+                        }
+                    };
+
+                    if let Err(e) = mgr.add_context(name, context) {
+                        error!("Failed to create context: {}", e);
+                        std::process::exit(1);
+                    }
+                    info!("Context '{}' created.", name);
+                }
+                Some(("rm", args)) => {
+                    let name = args.get_one::<String>("name").unwrap();
+                    if let Err(e) = mgr.remove_context(name) {
+                        error!("Failed to remove context: {}", e);
+                        std::process::exit(1);
+                    }
+                    info!("Context '{}' removed.", name);
+                }
+                Some(("ls", _)) => {
+                    let contexts = mgr.list_contexts();
+                    let current = mgr.current_name();
+                    for ctx in contexts {
+                        if Some(&ctx) == current.as_ref() {
+                            println!("* {}", ctx);
+                        } else {
+                            println!("  {}", ctx);
+                        }
+                    }
+                }
+                Some(("show", _)) => {
+                    if let Some(name) = mgr.current_name() {
+                        println!("{}", name);
+                    } else {
+                        println!("No context set (defaulting to local)");
+                    }
+                }
+                Some(("use", args)) => {
+                    let name = args.get_one::<String>("name").unwrap();
+                    if let Err(e) = mgr.set_current(name) {
+                        error!("Failed to set context: {}", e);
+                        std::process::exit(1);
+                    }
+                    info!("Switched to context '{}'.", name);
+                }
+                _ => unreachable!(),
+            }
+        }
         _ => unreachable!("Exhausted list of subcommands and subcommand_required prevents `None`"),
     }
 }