refractr/src/refractr.rs

/*
 * Copyright 2025 Bryson Steck <me@brysonsteck.xyz>
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

use crate::common::{self, ExitCode, ReturnData};
use crate::config::{Config, ConfigFile};

use git2::build::RepoBuilder;
use git2::string_array::StringArray;
use git2::{CertificateCheckStatus, Cred, FetchOptions, PushOptions, RemoteCallbacks, Repository};
use git2::{Error, ErrorCode};
use hex;
use sha2::{Digest, Sha256};
use std::env;
use std::fs;
use std::path::{Path, PathBuf, MAIN_SEPARATOR_STR};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;
use std::time;

pub struct Refractr {
  pub docker: bool,
  pub pid: u32,
  pub strict: bool,
  pub unix: bool,
  pub verbose: u8,
}

struct OpenedRepository {
  repo: Repository,
  path: String,
  remotes: Vec<String>,
  cfg: Config,
  ssh: bool,
}

impl Refractr {
  fn set_up_work_dir(&self, work_dir: PathBuf) -> Result<String, String> {
    if let Err(e) = fs::create_dir_all(&work_dir) {
      return Err(format!(
        "could not create working directory: {}: {}",
        work_dir.to_string_lossy().to_string(),
        e
      ));
    }
    Ok(work_dir.to_string_lossy().to_string())
  }

  fn set_up_ssh(&self, key_path: String, strict: bool) -> Result<RemoteCallbacks, String> {
    let mut cb = RemoteCallbacks::new();
    cb.credentials(move |_, _, _| Cred::ssh_key("git", None, Path::new(&key_path), None));
    cb.certificate_check(move |cert, url| {
      let sha256 = hex::encode(cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec());
      if strict {
        common::error(
          format!(
            "unknown host {} with sha256 host key {}, exiting",
            url, sha256
          ),
          ExitCode::RemoteError,
        );
      } else {
        common::warning(format!(
          "unknown host {} with sha256 host key {}, implicitly trusting",
          url, sha256
        ));
        common::warning(format!(
          "to suppress this warning in the future, add this host to your known_hosts file"
        ));
      }
      Ok(CertificateCheckStatus::CertificateOk)
    });
    Ok(cb)
  }

  fn get_refs(&self, branches: &Vec<String>, tags: Option<StringArray>) -> Vec<String> {
    let mut refs_branches = Vec::new();
    for branch in branches {
      refs_branches.push(format!("refs/heads/{}", branch));
    }
    if let Some(tags) = tags {
      for tag in &tags {
        refs_branches.push(format!("refs/tags/{}", tag.unwrap()))
      }
    }
    refs_branches
  }

  fn fast_forward(&self, repo_dir: &str, branches: &Vec<String>) -> Result<(), Error> {
    let repo = Repository::open(repo_dir)?;

    common::verbose(self.verbose, 2, format!("Pulling origin"));
    repo.find_remote("origin")?.fetch(&branches, None, None)?;

    for branch in branches {
      let refname = format!("refs/remotes/origin/{}", branch);
      let mut reference = repo.find_reference(&refname)?;
      reference.rename(format!("refs/heads/{}", branch).as_str(), true, "")?;
    }

    Ok(())
  }

  fn fetch(
    &self,
    repo: &Repository,
    branches: &Vec<String>,
    ssh: bool,
    ssh_key: &String,
    strict: bool,
  ) -> Result<(), Error> {
    let mut fo = FetchOptions::new();
    if ssh {
      match self.set_up_ssh(ssh_key.clone(), strict.clone()) {
        Ok(cb) => {
          fo.remote_callbacks(cb);
          ()
        },
        Err(e) => common::error(
          format!("error setting up ssh: {}", e),
          ExitCode::ConfigError,
        ),
      };
    }
    fo.download_tags(git2::AutotagOption::All);
    repo
      .find_remote("origin")?
      .fetch(&branches, Some(&mut fo), None)?;

    Ok(())
  }

  fn set_up_refs(&self, repo: &Repository, branches: &Vec<String>) -> Result<(), Error> {
    for branch in branches {
      let mut fetch_head =
        repo.find_reference(format!("refs/remotes/origin/{}", branch).as_str())?;
      fetch_head.rename(format!("refs/heads/{}", branch).as_str(), true, "")?;
    }

    Ok(())
  }

  fn make_remotes<'a>(
    &self,
    repo: &'a Repository,
    cfg: &ConfigFile,
  ) -> Result<Vec<String>, String> {
    // create remotes for each "to" repo
    let mut remote_list = Vec::new();
    for to in &cfg.config.to {
      let mut hasher = Sha256::new();
      hasher.update(to);
      let remote_id = format!("refractr-{}", &hex::encode(hasher.finalize())[..8]);
      common::verbose(
        self.verbose,
        2,
        format!("Attempting to create remote {} for url {}", remote_id, to),
      );
      match repo.remote(remote_id.as_str(), to) {
        Ok(_) => remote_list.push(remote_id),
        Err(e) => {
          if e.code() == ErrorCode::Exists {
            common::warning(format!("remote {} already exists, skipping", remote_id));
            remote_list.push(remote_id)
          } else {
            return Err(format!("failed to create remote: {}", e));
          }
        },
      }
    }

    Ok(remote_list)
  }

  fn push_remotes(
    &self,
    cfg: &Config,
    repo: &Repository,
    remote_list: &Vec<String>,
  ) -> Result<(), String> {
    for id in remote_list {
      let mut remote = repo.find_remote(&id).unwrap();
      common::verbose(
        self.verbose,
        1,
        format!("Pushing to remote: {}", remote.url().unwrap()),
      );
      let mut po = PushOptions::new();
      match self.set_up_ssh(cfg.git.ssh_identity_file.clone(), self.strict.clone()) {
        Ok(cb) => {
          po.remote_callbacks(cb);
          ()
        },
        Err(e) => common::error(
          format!("error setting up ssh: {}", e),
          ExitCode::ConfigError,
        ),
      };

      let mut refs = Vec::new();
      let mut refs_str = String::new();
      let strings = self.get_refs(
        &cfg.branches,
        match cfg.push_tags {
          true => Some(repo.tag_names(None).unwrap()),
          false => None,
        },
      );
      for branch in &strings {
        refs.push(branch.as_str());
        refs_str.push_str(format!("{} ", branch).as_str());
      }

      common::verbose(self.verbose, 4, format!("ref list: {}", refs_str));
      match remote.push::<&str>(&refs, Some(&mut po)) {
        Ok(_) => (),
        Err(e) => {
          if self.strict {
            return Err(format!(
              "failed to push to remote: {}: {}",
              remote.url().unwrap(),
              e
            ));
          } else {
            common::warning(format!(
              "failed to push to remote: {}: {}",
              remote.url().unwrap(),
              e
            ))
          }
        },
      }
    }

    Ok(())
  }

  fn looper(&self, repos: Vec<OpenedRepository>) -> Result<(), ReturnData> {
    let mut current_ints = Vec::new();
    let running = Arc::new(AtomicBool::new(true));
    let r = running.clone();
    let count = repos.len();
    for i in 0..repos.len() {
      current_ints.push(i64::from(
        repos[i].cfg.schedule.interval.unwrap().unsigned_abs(),
      ));
    }
    let original_ints = current_ints.clone();

    ctrlc::set_handler(move || {
      r.store(false, Ordering::SeqCst);
    })
    .expect("Failed to set ^C handler");

    common::verbose(self.verbose, 1, format!("Starting scheduled loop"));
    let mut do_break = false;
    while !do_break {
      do_break = true;
      let min = *current_ints.iter().min().unwrap();
      let sleep_int = time::Duration::from_secs(min as u64);
      let now = time::Instant::now();

      common::verbose(
        self.verbose,
        2,
        format!("Sleeping for {} seconds", sleep_int.as_secs()),
      );
      while running.load(Ordering::SeqCst) {
        thread::sleep(time::Duration::from_millis(200));
        if now.elapsed().as_secs() >= sleep_int.as_secs() {
          common::verbose(self.verbose, 3, format!("Thread has awoken!"));
          for i in 0..count {
            current_ints[i] -= now.elapsed().as_secs() as i64;
            common::verbose(
              self.verbose,
              4,
              format!("checking repo: {}", repos[i].cfg.from),
            );
            if current_ints[i] <= 0 {
              common::verbose(self.verbose, 4, format!("repo is ready for push"));
              common::verbose(
                self.verbose,
                2,
                format!("Interval for {} has arrived, pulling", repos[i].cfg.from),
              );
              if let Err(e) = self.fetch(
                &repos[i].repo,
                &repos[i].cfg.branches,
                repos[i].ssh,
                &repos[i].cfg.git.ssh_identity_file,
                self.strict.clone(),
              ) {
                common::error(
                  format!("failed to fetch repo {}: {}", repos[i].cfg.from, e),
                  ExitCode::FetchError,
                );
              }

              let _ = self.fast_forward(&repos[i].path, &repos[i].cfg.branches);
              if let Err(e) = self.push_remotes(&repos[i].cfg, &repos[i].repo, &repos[i].remotes) {
                common::error(e, ExitCode::PushError)
              };
              current_ints[i] = original_ints[i].clone();
            }
            common::verbose(
              self.verbose,
              4,
              format!("repo remaining time is now {}", current_ints[i]),
            );
          }
          do_break = false;
          break;
        }
      }
    }

    common::verbose(self.verbose, 1, format!("Exited looper due to ^C"));
    Ok(())
  }

  pub fn run(&self, cfgs: Vec<ConfigFile>) -> Result<(), ReturnData> {
    common::verbose(self.verbose, 3, format!("Starting main refractr loop"));
    let mut loop_repos = Vec::new();

    for cfg in cfgs {
      // set up the working directory
      common::verbose(self.verbose, 3, format!("Loading config: {}", cfg.path));
      let work_dir = self.set_up_work_dir(match &cfg.config.work_dir {
        None => {
          if cfg!(windows) {
            PathBuf::from(format!("{}\\refractr", env::var("TEMP").unwrap()))
          } else {
            PathBuf::from("/tmp/refractr")
          }
        },
        Some(path) => PathBuf::from(path),
      });
      let path_str = match work_dir {
        Ok(p) => p,
        Err(e) => {
          return Err(ReturnData {
            code: ExitCode::FilesystemError,
            msg: e,
          })
        },
      };
      common::verbose(
        self.verbose,
        2,
        format!("Created working directory: {}", &path_str),
      );
      let repo_name = match &cfg.config.from.split("/").last() {
        Some(split) => split.to_string(),
        None => {
          return Err(ReturnData {
            code: ExitCode::ParseError,
            msg: format!("failed to parse repository name"),
          })
        },
      };

      let ssh = cfg.config.from.starts_with("ssh://");
      let mut builder = RepoBuilder::new();
      let mut fo = FetchOptions::new();

      // make initial clone
      if ssh {
        match self.set_up_ssh(
          cfg.config.git.ssh_identity_file.clone(),
          self.strict.clone(),
        ) {
          Ok(cb) => {
            fo.remote_callbacks(cb);
            ()
          },
          Err(e) => common::error(
            format!("error setting up ssh: {}", e),
            ExitCode::ConfigError,
          ),
        };
      }
      fo.download_tags(git2::AutotagOption::All);
      builder.fetch_options(fo);

      let repo_dir = format!("{}{}{}", &path_str, MAIN_SEPARATOR_STR, repo_name);
      common::verbose(
        self.verbose,
        1,
        format!("Cloning repository {} to {}", &cfg.config.from, &repo_dir),
      );
      let repo = match builder.clone(&cfg.config.from, Path::new(&repo_dir)) {
        Ok(repo) => repo,
        Err(e) => {
          if e.code() != ErrorCode::Exists {
            common::error(
              format!("failed to clone repo to {}: {}", repo_dir, e),
              ExitCode::FilesystemError,
            );
          }
          common::warning(format!(
            "found existing repo at {}, attempting to use",
            repo_dir
          ));
          match self.fast_forward(&repo_dir, &cfg.config.branches) {
            Ok(_) => {
              if let Ok(repo) = Repository::open(Path::new(&repo_dir)) {
                repo
              } else {
                return Err(ReturnData {
                  code: ExitCode::RepositoryError,
                  msg: format!("failed to obtain existing repo"),
                });
              }
            },
            Err(e) => {
              return Err(ReturnData {
                code: ExitCode::RepositoryError,
                msg: format!("failed to obtain existing repo: {}", e),
              })
            },
          }
        },
      };

      if let Err(e) = self.set_up_refs(&repo, &cfg.config.branches) {
        common::error(
          format!("failed to set up refs: {}", e),
          ExitCode::RepositoryError,
        );
      }
      if let Err(e) = self.fetch(
        &repo,
        &cfg.config.branches,
        ssh,
        &cfg.config.git.ssh_identity_file,
        self.strict.clone(),
      ) {
        common::error(
          format!("failed to fetch repo {}: {}", cfg.config.from, e),
          ExitCode::FetchError,
        );
      }

      let remotes = match self.make_remotes(&repo, &cfg) {
        Ok(v) => v,
        Err(e) => {
          return Err(ReturnData {
            code: ExitCode::RemoteError,
            msg: e,
          })
        },
      };
      if let Err(e) = self.push_remotes(&cfg.config, &repo, &remotes) {
        common::error(e, ExitCode::PushError);
      }

      if cfg.config.schedule.enabled {
        loop_repos.push(OpenedRepository {
          repo,
          path: repo_dir,
          remotes,
          cfg: cfg.config,
          ssh,
        });
      }
    }

    if loop_repos.len() >= 1 {
      common::verbose(
        self.verbose,
        2,
        format!(
          "{} configs have schedules enabled, setting up looper",
          loop_repos.len()
        ),
      );
      return self.looper(loop_repos);
    } else {
      common::verbose(
        self.verbose,
        2,
        format!("No scheduled configs found, exiting refractr"),
      );
    }

    Ok(())
  }
}