ptth/crates/ptth_server/src/bin/ptth_server.rs

#![warn (clippy::pedantic)]

use std::{
	fs::File,
	path::{Path, PathBuf},
};

use structopt::StructOpt;

use ptth_server::{
	load_toml,
	prelude::*,
	run_server,
};

#[derive (Debug, StructOpt)]
struct Opt {
	#[structopt (long)]
	auto_gen_key: bool,
	
	#[structopt (long)]
	throttle_upload: bool,
	
	#[structopt (long)]
	file_server_root: Option <PathBuf>,
	
	#[structopt (long)]
	asset_root: Option <PathBuf>,
	
	#[structopt (long)]
	config_path: Option <PathBuf>,
	
	#[structopt (long)]
	name: Option <String>,
	
	#[structopt (long)]
	print_tripcode: bool,
	
	#[structopt (long)]
	relay_url: Option <String>,
}

#[derive (Default, serde::Deserialize)]
pub struct ConfigFile {
	pub name: Option <String>,
	pub api_key: String,
	pub relay_url: Option <String>,
	pub file_server_root: Option <PathBuf>,
}

fn gen_and_save_key (path: &Path) -> anyhow::Result <()> {
	let api_key = ptth_core::gen_key ();
	
	let mut f = File::create (path).with_context (|| format! ("Can't create config file `{:?}`", path))?;
	
	#[cfg (unix)]
	{
		use std::os::unix::fs::PermissionsExt;
		
		let metadata = f.metadata ()?;
		let mut permissions = metadata.permissions ();
		permissions.set_mode (0o600);
		f.set_permissions (permissions)?;
	}
	#[cfg (not (unix))]
	{
		tracing::warn! ("Error VR6VW5QT: API keys aren't protected from clients on non-Unix OSes yet");
	}
	
	f.write_all (format! ("api_key = \"{}\"\n", api_key).as_bytes ())?;
	
	Ok (())
}

#[tokio::main]
async fn main () -> Result <(), anyhow::Error> {
	tracing_subscriber::fmt::init ();
	
	let opt = Opt::from_args ();
	let asset_root = opt.asset_root;
	
	let path = opt.config_path.clone ().unwrap_or_else (|| PathBuf::from ("./config/ptth_server.toml"));
	
	let config_file: ConfigFile = if opt.auto_gen_key {
		// If we're in autonomous mode, try harder to fix things
		
		match load_toml::load (&path) {
			Err (_) => {
				gen_and_save_key (&path)?;
				
				load_toml::load (&path)?
			},
			Ok (x) => x,
		}
	}
	else {
		match load_toml::load (&path) {
			Err (ptth_server::errors::LoadTomlError::Io (_)) => bail! ("API key not provided in config file and auto-gen-key not provided"),
			Ok (x) => x,
			Err (e) => return Err (e.into ()),
		}
	};
	
	let config_file = ptth_server::ConfigFile {
		name: opt.name.or (config_file.name).ok_or (anyhow::anyhow! ("`name` must be provided in command line or config file"))?,
		api_key: config_file.api_key,
		relay_url: opt.relay_url.or (config_file.relay_url).ok_or (anyhow::anyhow! ("`--relay-url` must be provided in command line or `relay_url` in config file"))?,
		file_server_root: opt.file_server_root.or (config_file.file_server_root),
		throttle_upload: opt.throttle_upload,
	};
	
	if opt.print_tripcode {
		println! (r#"name = "{}""#, config_file.name);
		println! (r#"tripcode = "{}""#, config_file.tripcode ());
		return Ok (());
	}
	
	run_server (
		config_file, 
		ptth_core::graceful_shutdown::init (),
		Some (path),
		asset_root
	).await?;
	
	Ok (())
}