Initial commit

src/main.rs

@@ -0,0 +1,157 @@
+use std::io::Cursor;
+use std::{fs::File, io::BufReader};
+
+use log::*;
+use serde::{Deserialize, Serialize};
+use telegram_bot::CanReplySendMessage;
+use telegram_bot::{
+	Api, CanGetFile, CanReplySendAudio, InputFileUpload, MessageKind, UpdateKind, Voice,
+};
+use thiserror::Error;
+use tokio_stream::StreamExt;
+
+#[derive(Serialize, Deserialize, Debug)]
+struct Config {
+	token: String,
+}
+
+#[derive(Error, Debug)]
+enum UnresolvableError {
+	#[error(transparent)]
+	ConfigLoadError(ConfigLoadError),
+	#[error(transparent)]
+	ApiError(#[from] telegram_bot::Error),
+}
+
+#[derive(Error, Debug)]
+enum ConfigLoadError {
+	#[error("Failed to open config file: {0}")]
+	OpenFailed(#[source] std::io::Error),
+	#[error("Failed to read config file: {0}")]
+	ReadError(#[source] serde_json::Error),
+}
+
+#[derive(Error, Debug)]
+enum RuntimeError {
+	#[error("Couldn't get file information")]
+	FileInfoRequestFailed(#[source] telegram_bot::Error),
+	#[error("File has no URL")]
+	FileNoUrl,
+	#[error("Cannot download file")]
+	FileDownloadFailed(#[source] reqwest::Error),
+	#[error("Cannot decode audio file")]
+	FileDecodeError(#[source] ogg_opus::Error),
+	#[error("Cannot encode audio file")]
+	FileEncodeError(#[source] ogg_opus::Error),
+	#[error("Internal error")]
+	IOError(#[from] std::io::Error),
+}
+
+impl Config {
+	fn load() -> Result<Self, ConfigLoadError> {
+		let file = File::open("config.json").map_err(ConfigLoadError::OpenFailed)?;
+		let reader = BufReader::new(file);
+		serde_json::from_reader(reader).map_err(ConfigLoadError::ReadError)
+	}
+}
+
+#[tokio::main]
+async fn main() -> Result<(), UnresolvableError> {
+	pretty_env_logger::init();
+	info!("Starting amplifier-bot");
+
+	info!("Reading configuration");
+	let config = Config::load().map_err(UnresolvableError::ConfigLoadError)?;
+
+	info!("Long-polling for updates...");
+	let api = Api::new(&config.token);
+	let mut stream = api.stream();
+	while let Some(update) = stream.next().await {
+		// If the received update contains a new message...
+		let update = update?;
+		if let UpdateKind::Message(message) = update.kind {
+			if let MessageKind::Voice { data } = &message.kind {
+				let result = handle_voice_message(&config, &api, data).await;
+				match result {
+					Ok(audio) => {
+						api.send(
+							message.audio_reply(InputFileUpload::with_data(audio, "amplified.ogg")),
+						)
+						.await?;
+					}
+					Err(err) => {
+						error!("{:?}", err);
+						api.send(message.text_reply(format!("Error: {}", err)))
+							.await?;
+					}
+				}
+			}
+		}
+	}
+
+	Ok(())
+}
+
+async fn handle_voice_message(
+	config: &Config,
+	api: &Api,
+	data: &Voice,
+) -> Result<Vec<u8>, RuntimeError> {
+	info!("Amplifying voice message with duration {}s", data.duration);
+
+	let file_info = api
+		.send(data.get_file())
+		.await
+		.map_err(RuntimeError::FileInfoRequestFailed)?;
+
+	let url = file_info
+		.get_url(&config.token)
+		.ok_or(RuntimeError::FileNoUrl)?;
+
+	let ogg_data_in = reqwest::blocking::get(url)
+		.and_then(|response| response.bytes())
+		.map_err(RuntimeError::FileDownloadFailed)?;
+
+	let (audio, play_data) = ogg_opus::decode::<_, 48000>(Cursor::new(ogg_data_in))
+		.map_err(RuntimeError::FileDecodeError)?;
+
+	if play_data.channels != 1 {
+		warn!(
+			"Encountered unexpected voice file with {} channels",
+			play_data.channels
+		);
+	}
+
+	let mut last_audio = audio.clone();
+	let mut amplification_factor = 2;
+	loop {
+		let (current_audio, clipping_rate) = amplify_audio(&audio, amplification_factor);
+		if clipping_rate > 1.0 / (48000.0 * 10.0) {
+			break;
+		}
+		last_audio = current_audio;
+		amplification_factor += 1;
+	}
+
+	debug!("Amplified with factor {}", amplification_factor - 1);
+
+	let ogg_data_out =
+		ogg_opus::encode::<48000, 1>(&last_audio).map_err(RuntimeError::FileEncodeError)?;
+
+	Ok(ogg_data_out)
+}
+
+fn amplify_audio(original: &[i16], factor: i16) -> (Vec<i16>, f64) {
+	let mut modified_audio = original.to_owned();
+	let mut clipping_samples = 0;
+	for sample in modified_audio.iter_mut() {
+		*sample = sample.checked_mul(factor).unwrap_or_else(|| {
+			clipping_samples += 1;
+			sample.saturating_mul(factor)
+		});
+	}
+
+	let clipping_rate = clipping_samples as f64 / modified_audio.len() as f64;
+	info!("samples: {}    rate: {}", clipping_samples, clipping_rate);
+	(modified_audio, clipping_rate)
+}