initial commit

.gitignore

@@ -0,0 +1,22 @@
+# Generated by Cargo
+# will have compiled files and executables
+debug
+target
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb
+
+# Generated by cargo mutants
+# Contains mutation testing data
+**/mutants.out*/
+
+# RustRover
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+

Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+name = "atproto-ssh-tool"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+anyhow = "1.0.102"
+chrono = { version = "0.4.44", features = ["serde"] }
+clap = { version = "4.5.60", features = ["cargo", "derive"] }
+jacquard = "0.9.5"
+jacquard-api = "0.9.5"
+jacquard-common = "0.9.5"
+jacquard-derive = "0.9.5"
+jacquard-identity = "0.9.5"
+jacquard-lexgen = "0.9.5"
+jacquard-lexicon = "0.9.5"
+owo-colors = "4.3.0"
+reqwest = { version = "0.13.2", features = ["json"] }
+rpassword = "7.4.0"
+rustversion = "1.0.22"
+serde = { version = "1.0.228", features = ["derive"] }
+ssh-key = { version = "0.6.7", features = ["serde"] }
+tokio = { version = "1.50.0", features = ["full"] }
+unicode-segmentation = "1.12.0"
+
+[build-dependencies]
+jacquard-lexicon = "0.9.5"
+
+# So cargo shuts up
+[features]
+default = ["pink_grace"]
+pink_grace = []

build.rs

@@ -0,0 +1,16 @@
+use jacquard_lexicon::codegen::CodeGenerator;
+use jacquard_lexicon::corpus::LexiconCorpus;
+use std::env;
+use std::path::Path;
+
+fn main() {
+    let out_dir = env::var_os("OUT_DIR").unwrap();
+    let out_path = Path::new(&out_dir);
+
+    let corpus = LexiconCorpus::load_from_dir(Path::new("lexicons/")).unwrap();
+
+    CodeGenerator::new(&corpus, "pink_lexicon")
+        .write_to_disk(out_path)
+        .unwrap();
+    println!("cargo::rerun-if-changed=lexicons/");
+}

codebook.toml

@@ -0,0 +1,9 @@
+words = [
+    "akc",
+    "atproto",
+    "bsky",
+    "pds",
+    "sshd",
+    "usera",
+    "userb",
+]

lexicons/sshKey.json

@@ -0,0 +1,44 @@
+{
+  "defs": {
+    "main": {
+      "description": "Public SSH Key",
+      "key": "tid",
+      "record": {
+        "properties": {
+          "createdAt": {
+            "format": "datetime",
+            "type": "string"
+          },
+          "keyString": {
+            "maxLength": 3000,
+            "type": "string"
+          },
+          "name": {
+            "description": "Human name for SSH key",
+            "maxGraphemes": 64,
+            "maxLength": 640,
+            "minLength": 1,
+            "type": "string"
+          },
+          "tags": {
+            "type": "array",
+            "items": {
+              "type": "string",
+              "maxLength": 64
+            },
+            "maxLength": 10
+          }
+        },
+        "required": [
+          "name",
+          "keyString",
+          "createdAt"
+        ],
+        "type": "object"
+      },
+      "type": "record"
+    }
+  },
+  "id": "pink.grace.sshKey",
+  "lexicon": 1
+}

src/main.rs

@@ -0,0 +1,365 @@
+// atproto-ssh-tool by grace.pink
+
+// ignore warnings from generated code
+#[allow(warnings)]
+mod generated {
+    include!(concat!(env!("OUT_DIR"), "/lib.rs"));
+}
+use generated::*;
+
+use crate::pink_grace::ssh_key::SshKey;
+
+use std::fs;
+use std::path::PathBuf;
+use std::str::FromStr;
+
+use chrono::Utc;
+use clap::{Parser, Subcommand};
+use jacquard::api::com_atproto::repo::list_records::ListRecords;
+use jacquard::client::{Agent, AgentSessionExt, FileAuthStore, MemoryCredentialSession};
+use jacquard::oauth::client::OAuthClient;
+use jacquard::oauth::loopback::LoopbackConfig;
+use jacquard::types::ident::AtIdentifier;
+use jacquard::types::nsid::Nsid;
+use jacquard::types::string::Handle;
+use jacquard::types::value::from_data;
+use jacquard::{CowStr, IntoStatic, prelude::*};
+use owo_colors::OwoColorize;
+use reqwest::Url;
+
+/// Manage SSH keys with ATProto
+///
+/// Written by @grace.pink
+#[derive(Parser, Debug, Clone)]
+#[command(
+    author,
+    version,
+    about = "Manage SSH keys with ATProto",
+    after_help = "Examples:
+  atproto-ssh-tool put -a alice.bsky.social -n \"laptop\" -k ~/.ssh/id_ed25519.pub
+  atproto-ssh-tool get -a alice.bsky.social
+  atproto-ssh-tool akc -m alice=alice.bsky.social -u %u"
+)]
+struct Args {
+    /// OAuth Session Store
+    #[arg(long, default_value = "/tmp/jacquard-oauth-session.json")]
+    store: String,
+
+    /// Handle of ATProto user to interact with
+    #[arg(short, long, global = true)]
+    atproto_handle: Option<String>,
+
+    /// Password to authenticate. Leave blank to be prompted
+    #[arg(short, long, num_args = 0..=1, global = true)]
+    password: Option<Option<String>>,
+
+    /// Use Local OAuth. Requires network access
+    #[arg(short, long, global = true)]
+    oauth: bool,
+
+    #[command(subcommand)]
+    cmd: Commands,
+}
+
+#[derive(Subcommand, Debug, Clone)]
+enum Commands {
+    #[command(
+        about = "Get all SSH public keys for user",
+        after_help = "Examples:
+        atproto-ssh-tool get -a alice.bsky.social"
+    )]
+    Get {},
+    #[command(
+        about = "Upload a public SSH key to your ATProto PDS",
+        after_help = "Examples:
+        atproto-ssh-tool put -n \"work laptop\" -k ~/.ssh/id_ed25519.pub
+        atproto-ssh-tool put -n \"work laptop\" -k ~/.ssh/id_ed25519.pub -t work -t personal"
+    )]
+    Put {
+        /// human readable name
+        #[arg(short, long)]
+        name: String,
+
+        /// path to SSH public key
+        #[arg(short, long)]
+        key: PathBuf,
+
+        /// list of tags to add
+        #[arg(short, long)]
+        tags: Vec<CowStr<'static>>,
+    },
+    #[command(
+        aliases = ["akc"],
+        about = "Output select keys to stdout for use with AuthorizedKeyCommand",
+        after_help = "
+Note: In order to use this, you must use the AuthorizedKeyCommand and AuthorizedKeysCommandUser
+      options in your sshd_config. More details can be found in the sshd man page
+
+Examples:
+        AuthorizedKeyCommand atproto-ssh-tool akc -u %u -m root=alice.bsky.social
+        AuthorizedKeysCommandUser nobody
+
+        AuthorizedKeyCommand atproto-ssh-tool akc -u %u -m usera=alice.bsky.social -m userb=bob.bsky.social -i tag-a -e tag-b
+        AuthorizedKeysCommandUser nobody "
+    )]
+    AuthorizedKeyCommand {
+        /// The user attempting to log in. Will typically be %u
+        #[arg(short, long)]
+        username: String,
+
+        /// Mappings from unix user to ATProto handle. EX: user=handle.bsky.social
+        #[arg(long, short = 'm', value_parser = parse_key_val)]
+        user_mappings: Vec<(String, String)>,
+
+        /// Tags to include in the output. Will include all if empty
+        #[arg(short, long)]
+        include_tags: Vec<CowStr<'static>>,
+
+        /// Tags to remove in the output.
+        #[arg(short, long)]
+        exclude_tags: Vec<CowStr<'static>>,
+    },
+}
+
+fn parse_key_val(s: &str) -> Result<(String, String), String> {
+    let (a, b) = s
+        .split_once('=')
+        .ok_or_else(|| format!("expected `username=handle`, got `{s}`"))?;
+    Ok((a.to_string(), b.to_string()))
+}
+
+#[tokio::main]
+async fn main() {
+    let args = Args::parse();
+
+    if let Commands::AuthorizedKeyCommand {
+        user_mappings,
+        username,
+        include_tags,
+        exclude_tags,
+    } = &args.cmd
+    {
+        if let Err(e) =
+            authorized_key_command(username, user_mappings, include_tags, exclude_tags).await
+        {
+            eprintln!("{}: AuthorizedKeyCommand Failed", "ERROR".red());
+            eprintln!("{e:#}");
+        }
+        return;
+    }
+
+    // This clone is stupid but I do not particularly feel like fixing it rn
+    let handle = if let Some(hs) = args.clone().atproto_handle {
+        hs
+    } else {
+        println!(
+            "{}: Must provide an ATProto Handle with the -a flag",
+            "ERROR".red()
+        );
+        return;
+    };
+
+    if args.oauth {
+        let oauth = OAuthClient::with_default_config(FileAuthStore::new(&args.store));
+
+        let session = oauth
+            .login_with_local_server(&handle, Default::default(), LoopbackConfig::default())
+            .await
+            .unwrap();
+
+        let agent: Agent<_> = Agent::from(session);
+
+        dispatcher(args, &handle, agent).await;
+        return;
+    }
+
+    let pds_url = match get_pds_url(&handle).await {
+        Ok(url) => url,
+        Err(e) => {
+            eprintln!("{}: Failed to get PDS URL for user", "ERROR".red());
+            eprintln!("{e:#}");
+            return;
+        }
+    };
+    if let Some(opt_pass) = &args.password {
+        let password = if let Some(pass) = opt_pass {
+            pass.clone()
+        } else {
+            rpassword::prompt_password(format!("password for {handle}: ")).unwrap()
+        };
+
+        let session = match MemoryCredentialSession::authenticated(
+            handle.clone().into(),
+            password.into(),
+            None,
+            Some(pds_url),
+        )
+        .await
+        {
+            Ok(session) => session.0,
+            Err(e) => {
+                eprintln!("{}: Failed to authenticate", "ERROR".red());
+                eprintln!("{e:#}");
+                return;
+            }
+        };
+        let agent: Agent<_> = Agent::from(session);
+        dispatcher(args, &handle, agent).await;
+        return;
+    }
+    let agent = BasicClient::unauthenticated();
+    agent.set_base_uri(pds_url).await;
+
+    dispatcher(args, &handle, agent).await;
+}
+
+async fn get_pds_url(handle: &str) -> anyhow::Result<Url> {
+    let resolver = PublicResolver::default();
+    let did = resolver.resolve_handle(&Handle::new(handle)?).await?;
+    let doc_response = resolver.resolve_did_doc(&did).await?;
+    let doc = doc_response.parse()?;
+
+    doc.pds_endpoint()
+        .ok_or_else(|| anyhow::anyhow!("no PDS endpoint in DID doc"))
+}
+
+async fn dispatcher<T: AgentSessionExt>(args: Args, handle: &str, agent: Agent<T>) {
+    match args.cmd {
+        Commands::Put { name, key, tags } => put_key(agent, &name, &key, tags).await.unwrap(),
+        Commands::Get {} => print_keys(agent, handle).await.unwrap(),
+        Commands::AuthorizedKeyCommand { .. } => panic!("Impossible to reach this point"),
+    };
+}
+
+async fn put_key<T: AgentSessionExt>(
+    agent: Agent<T>,
+    name: &str,
+    key_path: &PathBuf,
+    tags: Vec<CowStr<'static>>,
+) -> anyhow::Result<()> {
+    let key_string = fs::read_to_string(key_path)?;
+
+    if key_string.contains("PRIVATE KEY") {
+        eprintln!(
+            "{}: Attempted to publish private key. Make sure to pass your public key instead",
+            "ERROR".red()
+        );
+
+        return Err(anyhow::anyhow!("Key file was private key"));
+    }
+
+    if key_string.trim().parse::<ssh_key::PublicKey>().is_err() {
+        eprintln!(
+            "{}: File passed is not a valid SSH public key",
+            "ERROR".red()
+        );
+
+        return Err(anyhow::anyhow!("Key file was invalid"));
+    }
+
+    let key = pink_grace::ssh_key::SshKey::new()
+        .name(name)
+        .key_string(key_string)
+        .tags(if tags.is_empty() { None } else { Some(tags) })
+        .created_at(Utc::now().fixed_offset())
+        .build();
+
+    if let Err(e) = agent.create_record(key, None).await {
+        eprintln!("{}: Failed to create ATProto record", "ERROR".red());
+        eprintln!("{e:#}");
+        return Err(anyhow::anyhow!("Failed to create ATProto record"));
+    }
+    println!("{}: Created SSH Key Record", "SUCCESS".green());
+    Ok(())
+}
+
+async fn get_key_list<'a, T: AgentSessionExt>(
+    agent: &'a Agent<T>,
+    handle: &'a str,
+) -> anyhow::Result<Vec<SshKey<'a>>> {
+    let request = ListRecords::new()
+        .repo(AtIdentifier::from_str(handle)?)
+        .collection(Nsid::from_str("pink.grace.sshKey")?)
+        .build();
+
+    let output = agent.send(request).await?.into_output()?;
+    let key_list: anyhow::Result<Vec<SshKey>> = output
+        .records
+        .iter()
+        .map(|record| {
+            from_data::<SshKey>(&record.value)
+                .map_err(anyhow::Error::from)
+                .map(|k| k.into_static())
+        })
+        .collect();
+    key_list
+}
+
+async fn print_keys<T: AgentSessionExt>(agent: Agent<T>, handle: &str) -> anyhow::Result<()> {
+    let keys = get_key_list(&agent, handle).await?;
+
+    println!("Keys for @{handle}");
+
+    for key in keys {
+        println!("name: {}", key.name);
+        println!(
+            "tags:  {}",
+            key.tags
+                .unwrap_or_default()
+                .iter()
+                .map(|s| { s.to_string() })
+                .collect::<Vec<String>>()
+                .join(", ")
+        );
+        println!("key:  {}", key.key_string);
+    }
+
+    Ok(())
+}
+
+async fn authorized_key_command(
+    username: &str,
+    user_mappings: &[(String, String)],
+    include_tags: &[CowStr<'static>],
+    exclude_tags: &[CowStr<'static>],
+) -> anyhow::Result<()> {
+    let handle = user_mappings
+        .iter()
+        .find(|(k, _)| k == username)
+        .map(|(_, v)| v)
+        .ok_or_else(|| anyhow::anyhow!("username has no associated handle"))?;
+
+    let pds_url = get_pds_url(handle).await?;
+
+    let agent = BasicClient::unauthenticated();
+    agent.set_base_uri(pds_url).await;
+
+    let keys = get_key_list(&agent, handle).await?;
+
+    let new_keys: Vec<SshKey<'_>> = keys
+        .into_iter()
+        .filter_map(|k| {
+            if let Some(tags) = &k.tags
+                && tags.iter().any(|t| exclude_tags.contains(t))
+            {
+                return None;
+            }
+            Some(k)
+        })
+        .filter_map(|k| {
+            if include_tags.is_empty() {
+                return Some(k);
+            }
+            match &k.tags {
+                Some(tags) if tags.iter().any(|t| include_tags.contains(t)) => Some(k),
+                _ => None,
+            }
+        })
+        .collect();
+
+    for key in new_keys {
+        println!("{}", key.key_string);
+    }
+
+    Ok(())
+}