add: rewriter logic and tests
This commit is contained in:
51
src/main.rs
Normal file
51
src/main.rs
Normal file
@@ -0,0 +1,51 @@
|
||||
use clap::Parser;
|
||||
|
||||
mod rewriter;
|
||||
|
||||
use env_logger::Env;
|
||||
use log::info;
|
||||
use rewriter::Rewriter;
|
||||
|
||||
/// Implements a HTTP server which allows clients to 'register'
|
||||
/// themselves. Their IP address will be used to replace a
|
||||
/// needle in a set of config files. This is intended to be
|
||||
/// used as a low-cost way of enabling Kubernetes ingress
|
||||
/// using nginx running on a machine that has a public port.
|
||||
///
|
||||
/// The needle is expected to be a dummy IP address; something
|
||||
/// fairly unique. The goal is to replace nginx files, where
|
||||
/// we often repeat lines if we want nginx to load balance between
|
||||
/// multiple destination IPs.
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(version, about, long_about = None)]
|
||||
struct Args {
|
||||
/// The needle that will be replaced. Anytime a line
|
||||
/// is encountered with this needle, the line is dropped
|
||||
/// and instead N lines (one per replacement) is added to
|
||||
/// the output.
|
||||
#[arg(short, long)]
|
||||
rewrite_string: String,
|
||||
|
||||
/// The folder which contains the templates that
|
||||
/// will be be searched for the needle.
|
||||
#[arg(short, long)]
|
||||
source_folder: String,
|
||||
|
||||
/// Where to write the replaced lines.
|
||||
#[arg(short, long)]
|
||||
dest_folder: String,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
// Log info and above by default
|
||||
env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
|
||||
let args = Args::parse();
|
||||
|
||||
let mut rewriter = Rewriter::new(args.rewrite_string);
|
||||
|
||||
rewriter.add_replacement(String::from("abc"));
|
||||
rewriter
|
||||
.rewrite_folder(&args.source_folder, &args.dest_folder)
|
||||
.unwrap();
|
||||
info!("Finished writing new config to {}", args.dest_folder);
|
||||
}
|
||||
117
src/rewriter.rs
Normal file
117
src/rewriter.rs
Normal file
@@ -0,0 +1,117 @@
|
||||
use anyhow::Result;
|
||||
use std::collections::HashSet;
|
||||
use std::path::Path;
|
||||
use std::{
|
||||
fs::{self, File},
|
||||
io::{BufReader, prelude::*},
|
||||
};
|
||||
|
||||
pub struct Rewriter {
|
||||
source: String,
|
||||
replacements: HashSet<String>,
|
||||
}
|
||||
|
||||
impl Rewriter {
|
||||
pub fn new(source: String) -> Self {
|
||||
Self {
|
||||
source,
|
||||
replacements: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_replacement(&mut self, replacement: String) {
|
||||
self.replacements.insert(replacement);
|
||||
}
|
||||
|
||||
pub fn remove_replacement(&mut self, replacement: &str) {
|
||||
self.replacements.remove(replacement);
|
||||
}
|
||||
|
||||
pub fn rewrite_folder(&self, src: &str, dst: &str) -> Result<()> {
|
||||
// Make sure we are deterministic; construct a list of strings and sort
|
||||
// them
|
||||
let mut replacements: Vec<String> = self.replacements.iter().map(|s| s.clone()).collect();
|
||||
replacements.sort();
|
||||
let dst_base = Path::new(dst);
|
||||
let mut to_visit = vec![String::from(src)];
|
||||
while to_visit.len() > 0 {
|
||||
let dir = to_visit.pop().unwrap();
|
||||
for entry in fs::read_dir(&dir)? {
|
||||
let entry = entry?;
|
||||
let metadata = entry.metadata()?;
|
||||
|
||||
// We need to find the destination this should be written into
|
||||
// First, calculate the 'relative' path after we trim the
|
||||
// the src prefix, then join that with the dst prefix
|
||||
let src_path = entry.path();
|
||||
let src_part_path = src_path.strip_prefix(&src)?;
|
||||
let dst_path = dst_base.join(src_part_path);
|
||||
|
||||
if metadata.is_dir() {
|
||||
// Create the directory, then carry on. Note that we explore the
|
||||
// src_path after creating dst_path.
|
||||
fs::create_dir(&dst_path)?;
|
||||
to_visit.push(src_path.into_os_string().into_string().unwrap());
|
||||
continue;
|
||||
}
|
||||
|
||||
// Open 2 files; one to read and translate, and one to write.
|
||||
let source_file = File::open(&src_path)?;
|
||||
let mut dest_file = File::create(&dst_path)?;
|
||||
let reader = BufReader::new(source_file);
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
// If the line is not subject to replacement, copy it and
|
||||
// carry on.
|
||||
if !line.contains(&self.source) {
|
||||
writeln!(dest_file, "{}", line)?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Else, repeat the line multiple times, replacing the string
|
||||
// in question
|
||||
for replacement in &replacements {
|
||||
let new_line = line.replace(&self.source, &replacement);
|
||||
writeln!(dest_file, "{}", new_line)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use include_directory::{Dir, include_directory};
|
||||
use tempdir::TempDir;
|
||||
|
||||
use super::Rewriter;
|
||||
|
||||
static TEST_FILES: Dir<'_> = include_directory!("$CARGO_MANIFEST_DIR/src/testdata");
|
||||
#[test]
|
||||
fn basic_test() {
|
||||
let testdata = TempDir::new("").unwrap();
|
||||
TEST_FILES.extract(testdata.path()).unwrap();
|
||||
let src = testdata.path().join("testsrc");
|
||||
let dst = TempDir::new("").unwrap();
|
||||
|
||||
let mut rewriter = Rewriter::new("to_be_replaced".into());
|
||||
rewriter.add_replacement("abc".into());
|
||||
rewriter.add_replacement("def".into());
|
||||
rewriter.add_replacement("zyx".into());
|
||||
rewriter.remove_replacement("zyx");
|
||||
rewriter
|
||||
.rewrite_folder(
|
||||
src.as_os_str().to_str().unwrap(),
|
||||
dst.path().as_os_str().to_str().unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Validate that everything matches
|
||||
assert!(
|
||||
dir_diff::is_different(testdata.path().join("testdst"), dst.path()).unwrap() == false
|
||||
);
|
||||
}
|
||||
}
|
||||
5
src/testdata/testdst/hello.txt
vendored
Normal file
5
src/testdata/testdst/hello.txt
vendored
Normal file
@@ -0,0 +1,5 @@
|
||||
This is a line
|
||||
This is abc line
|
||||
This is def line
|
||||
|
||||
This is another line
|
||||
4
src/testdata/testdst/recursive/world.txt
vendored
Normal file
4
src/testdata/testdst/recursive/world.txt
vendored
Normal file
@@ -0,0 +1,4 @@
|
||||
This is a abc line.
|
||||
This is a def line.
|
||||
|
||||
In a nested directory.
|
||||
4
src/testdata/testsrc/hello.txt
vendored
Normal file
4
src/testdata/testsrc/hello.txt
vendored
Normal file
@@ -0,0 +1,4 @@
|
||||
This is a line
|
||||
This is to_be_replaced line
|
||||
|
||||
This is another line
|
||||
3
src/testdata/testsrc/recursive/world.txt
vendored
Normal file
3
src/testdata/testsrc/recursive/world.txt
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
This is a to_be_replaced line.
|
||||
|
||||
In a nested directory.
|
||||
Reference in New Issue
Block a user