atcities.dev / upload / src / main.rs

use clap::{ArgAction, Parser};
use jacquard::api::com_atproto::repo::apply_writes::{
  self, ApplyWrites, ApplyWritesOutput, ApplyWritesWritesItem,
};
use jacquard::client::AgentSessionExt;
use jacquard::client::MemorySessionStore;
use jacquard::oauth::loopback::LoopbackConfig;
use jacquard::oauth::types::AuthorizeOptions;
use jacquard::types::string::{AtStrError, RecordKey, Rkey};
use jacquard::{AuthorizationToken, CowStr, atproto, oauth};
use jacquard::{
  Data,
  api::com_atproto::{self, repo::list_records::ListRecords},
  client::{Agent, credential_session::CredentialSession},
  cowstr::ToCowStr,
  identity::JacquardResolver,
  types::{ident::AtIdentifier, nsid::Nsid, string::AtprotoStr, uri::Uri},
  xrpc::XrpcExt,
};
use miette::{Context, IntoDiagnostic, Result};
use std::io::Write;
use std::{collections::HashMap, fs, path::PathBuf};

use crate::sitemap::{BlobRef, Sitemap, SitemapNode};

mod sitemap;
mod utils;

#[derive(Parser, Debug, Clone)]
#[command(version, about, long_about = None)]
struct Config {
  /// Handle or DID to authenticate
  #[arg(verbatim_doc_comment)]
  user: String,

  /// App password to authenticate the client
  /// Normal passwords also work but are not advised
  /// If ommited, oauth will be used instead
  /// Oauth is reccomended where possible.
  #[arg(verbatim_doc_comment, short = 'p', long = "password")]
  password: Option<String>,

  /// Include dotfiles in upload
  /// Default: false
  #[arg(verbatim_doc_comment, short = 'a', long = "all")]
  all_files: bool,

  /// Respect gitignore files
  /// Note: gitignore files are not uploaded unless --all is set
  /// Default: true
  #[arg(verbatim_doc_comment, short = 'g', long = "no-gitignore", action = ArgAction::SetFalse, )]
  git_ignore: bool,

  /// Directory to upload
  #[arg(verbatim_doc_comment)]
  dir: PathBuf,
}

async fn live_records(agent: &impl AgentSessionExt, config: Config) -> Result<Vec<String>> {
  // find live site records
  let mut cursor = None;
  let mut remote_records = Vec::new();
  let user = config.user.clone();
  let user = if user.contains(":") {
    AtIdentifier::Did(user.into())
  } else {
    AtIdentifier::Handle(user.into())
  };
  loop {
    let req = com_atproto::repo::list_records::ListRecords::new()
      .collection(
        Nsid::new("dev.atcities.route").expect("failed to generate dev.atcities.route nsid"),
      )
      .repo(user.clone())
      .limit(100)
      .maybe_cursor(cursor)
      .build();

    let res = agent
      .xrpc(agent.endpoint().await)
      .send::<ListRecords>(&req)
      .await?
      .into_output()?;

    for record in res.records {
      match record {
        Data::Object(obj) => {
          let obj = obj.0.clone();
          let uri = obj.get_key_value("uri").and_then(|x| match x.1 {
            Data::String(str) => Some(str),
            _ => None,
          });

          if let Some(uri) = uri
            && let AtprotoStr::Uri(uri) = uri
            && let Uri::At(uri) = uri
            && let Some(rkey) = uri.rkey()
          {
            let rkey = rkey.0.to_cowstr().to_string();
            remote_records.push(rkey);
          } else {
            panic!("Warning: pds returned invalid data.")
          }
        }
        _ => {
          panic!("Warning: pds returned invalid data.")
        }
      }
    }

    cursor = res.cursor;
    if cursor.is_none() {
      break;
    }
  }
  Ok(remote_records)
}

async fn upload_site_blobs(
  agent: &impl AgentSessionExt,
  _config: Config,
  local_sitemap: Sitemap,
) -> Result<Sitemap> {
  // upload local site blobs
  let mut new_sitemap: Sitemap = HashMap::new();
  for (k, v) in local_sitemap {
    print!("Uploading {k}... ");
    let _ = std::io::stdout().flush();
    let blob = match v.blob {
      BlobRef::Local(path) => path,
      BlobRef::Remote(_) => {
        panic!("Impossible state")
      }
    };
    let blob = fs::read(blob).into_diagnostic()?;

    let req = com_atproto::repo::upload_blob::UploadBlob::new()
      .body(blob.into())
      .build();
    let res = agent.send(req).await?.into_output()?;

    new_sitemap.insert(
      k,
      SitemapNode {
        mime_type: v.mime_type,
        blob: BlobRef::Remote(res.blob.into()),
      },
    );

    println!("Done!");
  }

  Ok(new_sitemap)
}

async fn update_remote_site(
  agent: &impl AgentSessionExt,
  config: Config,
  auth: AuthorizationToken<'_>,
  remote_records: Vec<String>,
  new_sitemap: Sitemap,
) -> Result<ApplyWritesOutput<'static>> {
  // batch delete/upload records
  let mut writes = Vec::new();
  let mut delete_records = remote_records
    .into_iter()
    .map(|x| {
      let rkey = RecordKey(Rkey::new_owned(x)?);
      Ok(ApplyWritesWritesItem::Delete(Box::new(
        apply_writes::Delete::builder()
          .collection(Nsid::raw("dev.atcities.route"))
          .rkey(rkey)
          .build(),
      )))
    })
    .collect::<Result<Vec<ApplyWritesWritesItem<'_>>, AtStrError>>()
    .into_diagnostic()?;

  let mut create_records = new_sitemap
    .into_iter()
    .map(|(k, v)| {
      let k = match k.as_str() {
        "404.html" => String::from("404"),
        "index.html" => String::from("/"),
        _ => match k.strip_suffix("/index.html") {
          Some(k) => format!("/{k}/"),
          None => format!("/{k}"),
        }
      };
      let rkey =
        utils::url_to_rkey(k).wrap_err("Invalid file path. Could not be converted to rkey")?;
      let rkey = RecordKey(Rkey::new_owned(rkey).into_diagnostic()?);
      let blob = match v.blob {
        BlobRef::Local(_) => panic!("Illegal local blob"),
        BlobRef::Remote(cid) => cid,
      };
      let data = atproto!({
        "page": {
          "$type": "dev.atcities.route#blob",
          "blob": {
            "$type": "blob",
            "ref": {
              "$link": blob.r#ref.as_str()
            },
            "mimeType": blob.mime_type.0.as_str(),
            "size": blob.size
          }
        }
      });
      Ok(ApplyWritesWritesItem::Create(Box::new(
        apply_writes::Create::builder()
          .collection(Nsid::raw("dev.atcities.route"))
          .rkey(rkey)
          .value(data)
          .build(),
      )))
    })
    .collect::<Result<Vec<ApplyWritesWritesItem<'_>>, miette::Error>>()?;

  writes.append(&mut delete_records);
  writes.append(&mut create_records);

  let repo = if config.user.contains(":") {
    AtIdentifier::Did(config.user.into())
  } else {
    AtIdentifier::Handle(config.user.into())
  };

  let req = com_atproto::repo::apply_writes::ApplyWrites::new()
    .repo(repo)
    .writes(writes)
    .build();

  let res = agent
    .xrpc(agent.endpoint().await)
    .auth(auth)
    .send::<ApplyWrites>(&req)
    .await?
    .into_output()?;

  Ok(res)
}

#[tokio::main]
async fn main() -> Result<(), miette::Error> {
  env_logger::init();
  // get config items
  let config = Config::parse();

  // get local site info
  let local_sitemap =
    sitemap::local_sitemap(config.dir.clone(), config.all_files, config.git_ignore)?;

  // create session
  if let Some(password) = config.password.clone() {
    let password = password.into();
    let client = JacquardResolver::default();
    let store = MemorySessionStore::default();
    let session = CredentialSession::new(store.into(), client.into());

    let auth = session
      .login(config.user.clone().into(), password, None, None, None)
      .await?;

    let agent = Agent::from(session);

    let remote_sitemap = live_records(&agent, config.clone()).await?;
    let new_sitemap = upload_site_blobs(&agent, config.clone(), local_sitemap).await?;
    let _ = update_remote_site(
      &agent,
      config.clone(),
      AuthorizationToken::Bearer(auth.access_jwt),
      remote_sitemap,
      new_sitemap,
    )
    .await?;

    println!(
      "Site is now updated. Live at {}",
      utils::site_handle(config.user)
    );
  } else {
    let oauth = oauth::client::OAuthClient::with_memory_store();
    let session = oauth
      .login_with_local_server(
        config.user.clone(),
        AuthorizeOptions::default(),
        LoopbackConfig::default(),
      )
      .await?;

    // sick and twisted reference mangling BUT it works So
    // tldr: the cowstr is a borrowed cowstr iiuc,
    // so it needs to be turned into an owned cowstr
    // to break reference to session which gets moved
    let auth = session.access_token().await;
    let auth = match auth {
      AuthorizationToken::Bearer(cow_str) => CowStr::copy_from_str(cow_str.as_str()),
      AuthorizationToken::Dpop(cow_str) => CowStr::copy_from_str(cow_str.as_str()),
    };

    println!("{}", auth);

    let agent = Agent::from(session);

    let remote_sitemap = live_records(&agent, config.clone()).await?;
    let new_sitemap = upload_site_blobs(&agent, config.clone(), local_sitemap).await?;
    let _ = update_remote_site(
      &agent,
      config.clone(),
      AuthorizationToken::Dpop(auth),
      remote_sitemap,
      new_sitemap,
    )
    .await?;

    println!(
      "Site is now updated. Live at {}",
      utils::site_handle(config.user)
    );
  };

  Ok(())
}