refractr/src/refractr.rs

use crate::freak_out;
use crate::common;
use crate::config::{Config, ConfigFile};

use git2::{FetchOptions, CertificateCheckStatus, Cred, PushOptions, RemoteCallbacks, Repository};
use git2::{Error, ErrorCode};
use git2::build::{RepoBuilder, CheckoutBuilder};
use hex;
use sha2::{Sha256, Digest};
use std::env;
use std::fs;
use std::path::{Path, PathBuf, MAIN_SEPARATOR_STR};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
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,
}

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

  fn get_refs(&self, branches: &Vec<String>) -> Vec<String> {
    let mut refs_branches = Vec::new();
    for branch in branches {
      refs_branches.push(format!("refs/heads/{}", branch));
    }
    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: Option<&String>) -> Result<(), Error> {
    match ssh {
      Some(key) => {
        let mut cb = RemoteCallbacks::new();
        let mut fo = FetchOptions::new();
        cb.credentials(move |_,_,_| Cred::ssh_key(
          "git",
          None,
          Path::new(&key),
          None));
        cb.certificate_check(|cert, url| {
          let mut sha256 = String::new();
          for i in cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec() {
            sha256.push_str(&hex::encode(i.to_string()));
          }
          common::warning(
            format!("implicitly trusting unknown host {} with sha256 host key {}",
            url,
            hex::encode(cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec())));
          common::warning(
            format!("to ignore this error in the future, add this host to your known_hosts file"));
          Ok(CertificateCheckStatus::CertificateOk)
        });
        fo.download_tags(git2::AutotagOption::All);
        fo.remote_callbacks(cb);
        repo.find_remote("origin")?.fetch(&branches, Some(&mut fo), None)?;
      },
      None => repo.find_remote("origin")?.fetch(&branches, None, 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 {
            freak_out!(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 callbacks = RemoteCallbacks::new();
      callbacks.credentials(|_,_,_| Cred::ssh_key(
        "git",
        None,
        &Path::new(&cfg.git.ssh_identity_file),
        None));
      callbacks.certificate_check(|cert, url| {
        let mut sha256 = String::new();
        for i in cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec() {
          sha256.push_str(&hex::encode(i.to_string()));
        }
        common::warning(
          format!("implicitly trusting unknown host {} with sha256 host key {}",
          url,
          hex::encode(cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec())));
        common::warning(
          format!("to ignore this error in the future, add this host to your known_hosts file"));
        Ok(CertificateCheckStatus::CertificateOk)
      });
      let mut push_options = PushOptions::new();
      push_options.remote_callbacks(callbacks);

      let mut refs = Vec::new();
      let strings = self.get_refs(&cfg.branches);
      for branch in &strings {
        refs.push(branch.as_str());
      }

      match remote.push::<&str>(&refs, Some(&mut push_options)) {
        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<(), String> {
    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(u64::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 min = *current_ints.iter().min().unwrap();
    let mut do_break = false;
    while !do_break {
      do_break = true;
      let sleep_int = time::Duration::from_secs(min);
      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_secs(1));
        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();
            if i <= 0 {
              current_ints[i] = original_ints[i].clone();
              common::verbose(
                self.verbose,
                2,
                format!("Interval for {} has arrived, pulling", repos[i].cfg.from));

              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) {
                  return Err(e)
                }
            }
          }
          do_break = false;
          break
        }
      }
    }

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

  pub fn run(&self, cfgs: Vec<ConfigFile>) -> Result<(), String> {
    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(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 => freak_out!(format!("failed to parse repository name"))
      };

      // make initial clone
      common::verbose(
        self.verbose,
        1,
        format!("Cloning repository: {}", &cfg.config.from));

      let mut builder = RepoBuilder::new();
      let mut cb = RemoteCallbacks::new();
      let mut fo = FetchOptions::new();
      
      if cfg.config.from.starts_with("ssh://") {
        let key_string = cfg.config.git.ssh_identity_file.clone();
        cb.credentials(move |_,_,_| Cred::ssh_key(
          "git",
          None,
          Path::new(&key_string),
          None));
        cb.certificate_check(|cert, url| {
          let mut sha256 = String::new();
          for i in cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec() {
            sha256.push_str(&hex::encode(i.to_string()));
          }
          common::warning(
            format!("implicitly trusting unknown host {} with sha256 host key {}",
            url,
            hex::encode(cert.as_hostkey().unwrap().hash_sha256().unwrap().to_vec())));
          common::warning(
            format!("to ignore this error in the future, add this host to your known_hosts file"));
          Ok(CertificateCheckStatus::CertificateOk)
        });
      }
      fo.download_tags(git2::AutotagOption::All);
      fo.remote_callbacks(cb);
      builder.fetch_options(fo);

      let repo_dir = format!("{}{}{}", &path_str, MAIN_SEPARATOR_STR, repo_name);
      let repo = match builder.clone(&cfg.config.from, Path::new(&repo_dir)) {
        Ok(repo) => repo,
        Err(_) => {
          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 {
              freak_out!(format!("failed to obtain existing repo"))
            },
            Err(e) => freak_out!(format!("failed to obtain existing repo: {}", e))
          }
        }
      };

      self.set_up_refs(&repo, &cfg.config.branches).unwrap();
      self.fetch(&repo, &cfg.config.branches, Some(&cfg.config.git.ssh_identity_file)).unwrap();

      let remotes = match self.make_remotes(&repo, &cfg) {
        Ok(v) => v,
        Err(e) => return Err(e)
      };
      if let Err(e) = self.push_remotes(&cfg.config, &repo, &remotes) {
        return Err(e)
      }

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

    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(())
  }
}