main.rs

use clap::Parser;
use lol_html::html_content::ContentType;
use lol_html::{element, rewrite_str, text, RewriteStrSettings};
use serde::{Deserialize, Serialize};
use std::borrow::{Borrow, BorrowMut};
use std::cell::RefCell;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Read, Write};
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use thiserror::Error;

/// Pipelined HTML ops
#[derive(Clone, Debug, Parser)]
#[clap(author, version, about)]
pub struct CliArgs {
	/// Path to the input html file
	#[clap(short, long, value_parser)]
	input: String,
	/// Path to write the transformed html to. If not present and --inline is not specified, the output will be written to stdout
	#[clap(short, long, value_parser)]
	output: Option<String>,
	/// Path to the file defining operations to perform on the html
	#[clap(short = 'f', long, value_parser)]
	opfile: String,
	/// Write the transformed html back to the original file. This will cause --output to be ignored
	#[clap(short = 'l', long, value_parser, default_value_t = false)]
	inline: bool,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "action")]
pub enum ConfigEntry {
	#[serde(rename = "wrap")]
	Wrap { template: String },
}

pub type ConfigFile = HashMap<String, ConfigEntry>;

#[derive(Error, Debug)]
pub enum AppError {
	#[error("Failed to open file: {0}")]
	FailedToOpen(String),
	#[error("Failed to read file: {0}")]
	FailedToRead(String),
	#[error("{0}")]
	IoError(#[from] std::io::Error),
	#[error("Opfile was not valid: {0}")]
	InvalidOpfile(toml::de::Error),
	#[error("Failed to match selector {0}")]
	SelectorFailed(String),
	#[error("Failed to apply edits {0}")]
	EditsFailed(#[from] lol_html::errors::RewritingError),
}

pub type AppResult<T> = Result<T, AppError>;

pub fn read_and_drop(path: &String) -> AppResult<String> {
	let mut file_handle = File::open(path).map_err(|_| AppError::FailedToOpen(path.clone()))?;
	let length = file_handle.metadata()?.len();
	let mut file_buffer = String::with_capacity(length as usize);
	file_handle
		.read_to_string(&mut file_buffer)
		.map_err(|_| AppError::FailedToRead(path.clone()))?;

	Ok(file_buffer)
}

fn main() -> AppResult<()> {
	let args: CliArgs = CliArgs::parse();
	let input_data = read_and_drop(&args.input)?;
	let opfile_data = read_and_drop(&args.opfile)?;

	let ops: ConfigFile = toml::from_str(opfile_data.as_str()).map_err(AppError::InvalidOpfile)?;

	let html = rewrite_str(
		input_data.as_str(),
		RewriteStrSettings {
			element_content_handlers: ops
				.into_iter()
				.flat_map(|(selector, op)| match op {
					ConfigEntry::Wrap { template } => {
						let mut tag_name = Arc::new(Mutex::new(String::new()));
						let mut attributes = Arc::new(Mutex::new(String::new()));
						let mut content = Arc::new(Mutex::new(String::new()));

						let el_content = content.clone();
						vec![
							element!(selector, move |e| {
								tag_name.lock().unwrap().push_str(&e.tag_name());

								let attr_list = e
									.attributes()
									.iter()
									.map(|a| format!(r#"{}="{}""#, a.name(), a.value()))
									.collect::<Vec<String>>()
									.join(" ");

								attributes.lock().unwrap().push_str(&attr_list);
								e.remove();

								let inner_template = template.clone();
								let inner_tag_name = tag_name.clone();
								let inner_attr = attributes.clone();
								let inner_content = el_content.clone();
								e.on_end_tag(move |end_tag| {
									let mut value = inner_template.replace(
										"{{{TAG}}}",
										inner_tag_name.clone().lock().unwrap().as_str(),
									);
									let mut value = value
										.replace("{{{ATTR}}}", inner_attr.lock().unwrap().as_str());
									let mut value = value.replace(
										"{{{CONTENT}}}",
										inner_content.lock().unwrap().as_str(),
									);

									end_tag.after(&value, ContentType::Html);

									end_tag.remove();
									Ok(())
								})
								.expect("Failed to handle end tag");
								Ok(())
							}),
							text!(selector, move |t| {
								content.lock().unwrap().push_str(t.as_str());
								t.remove();
								Ok(())
							}),
						]
					}
				})
				.collect(),
			document_content_handlers: vec![],
			strict: true,
			enable_esi_tags: true,
		},
	);

	let write_path = if args.inline {
		Some(args.input.clone())
	} else {
		args.output
	};

	if let Some(path) = write_path {
		File::create(&path)
			.map_err(|_| AppError::FailedToOpen(args.input.clone()))?
			.write_all(html?.as_bytes())
			.map_err(|_| AppError::FailedToRead(args.input.clone()))?;
	} else {
		println!("{}", html?);
	}

	Ok(())
}