use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use anyhow::Context as _;
use base64::Engine as _;
use knot_index::{Index, Resolved};
use knot_runtime::ReqwestHttp;
use knot_types::{AccountDid, OfferedKey};
use crate::config::ValidatedConfig;
use crate::datastore::record::Held;
use crate::datastore::{Datastore, schema};
use crate::out::out;
pub(crate) type Keys = BTreeMap<AccountDid, Vec<String>>;
fn load(config: &ValidatedConfig) -> anyhow::Result<Keys> {
let store = Datastore::at(&config.output_root);
Ok(store
.record::<schema::SshKeys>(&())
.read()?
.unwrap_or_default())
}
fn index(config: &ValidatedConfig) -> anyhow::Result<Index> {
let layout = config.layout()?;
let index = Index::new(&layout.meta_path(&config.knot_did)?, layout);
index.rebuild().context("rebuild registry index")?;
Ok(index)
}
pub(crate) fn accounts(config: &ValidatedConfig) -> anyhow::Result<BTreeSet<AccountDid>> {
pushers(config, &index(config)?)
}
fn pushers(config: &ValidatedConfig, index: &Index) -> anyhow::Result<BTreeSet<AccountDid>> {
let mut accounts = BTreeSet::from([config.owner.clone()]);
for repo in index.hosted_repos() {
if let Resolved::Ready(Some(owner)) = index.owner_of(&repo) {
accounts.insert(AccountDid::from(owner));
}
index.ensure_collaborators(&repo)?;
accounts.extend(index.effective_collaborators(&repo).ready_or_default());
}
Ok(accounts)
}
pub(crate) fn line_of(key: &OfferedKey) -> Option<String> {
let blob = key.as_bytes();
let length = u32::from_be_bytes(blob.get(..4)?.try_into().ok()?) as usize;
let algorithm = std::str::from_utf8(blob.get(4..4 + length)?).ok()?;
let encoded = base64::engine::general_purpose::STANDARD.encode(blob);
is_key((algorithm, &encoded)).then(|| format!("{algorithm} {encoded}"))
}
pub async fn sync(config: &Arc<ValidatedConfig>, all: bool) -> anyhow::Result<()> {
// hold the record until commit, and a parallel sync can't write over this one.
let store = Datastore::at(&config.output_root);
let record = store.record::<schema::SshKeys>(&());
let mut held = record.hold()?;
let stored = held.value();
let wanted = pushers(config, &index(config)?)?;
stored.retain(|account, _| wanted.contains(account));
let unread: Vec<_> = wanted
.iter()
.filter(|account| all || !stored.contains_key(*account))
.collect();
if unread.is_empty() {
return commit(held);
}
let atproto =
config.atproto(ReqwestHttp::new(config.http_limits()).context("init http transport")?)?;
let mut failed = Vec::new();
for account in unread {
match atproto.resolve_pubkeys(account).await {
Ok(keys) => {
stored.insert(account.clone(), keys.iter().filter_map(line_of).collect());
}
Err(error) if error.is_gone() => {
stored.insert(account.clone(), Vec::new());
}
Err(error) => failed.push(format!("{account}: {error}")),
}
}
commit(held)?;
anyhow::ensure!(
failed.is_empty(),
"couldn't read the keys of {}",
failed.join(", ")
);
Ok(())
}
fn commit(mut held: Held<'_, schema::SshKeys>) -> anyhow::Result<()> {
let keys = held.value();
let (accounts, count) = (keys.len(), keys.values().map(Vec::len).sum::<usize>());
held.commit()?;
out!("{accounts} accounts, {count} keys");
Ok(())
}
pub fn list(config: &ValidatedConfig, command: &str) -> anyhow::Result<String> {
anyhow::ensure!(
!command.contains(['"', '\\']) && !command.contains(char::is_control),
"the forced command can't contain a quote, a backslash or a control character"
);
let mut owners: BTreeMap<&str, BTreeSet<&AccountDid>> = BTreeMap::new();
let stored = load(config)?;
for (account, keys) in &stored {
for key in keys
.iter()
.filter(|key| key.split_once(' ').is_some_and(is_key))
{
owners.entry(key).or_default().insert(account);
}
}
Ok(owners
.into_iter()
.map(|(key, accounts)| {
let actors: String = accounts
.iter()
.map(|account| format!(" --actor {account}"))
.collect();
format!("command=\"{command}{actors}\",restrict {key}\n")
})
.collect())
}
/// a stored line with a line break or quote would forge an authorized_keys line
fn is_key((key_type, key_base64): (&str, &str)) -> bool {
["ssh-", "ecdsa-", "sk-"]
.iter()
.any(|prefix| key_type.starts_with(prefix))
&& key_type
.chars()
.all(|c| c.is_ascii_alphanumeric() || "-@.".contains(c))
&& !key_base64.is_empty()
&& base64::engine::general_purpose::STANDARD
.decode(key_base64)
.is_ok()
}