xenofem
/
screencap-bot
1
0
Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

Oneshot mode and general improved error handling

main
xenofem 2023-07-04 13:28:45 -04:00
parent 31e1b6302a
commit 8f74af7c7b
7 changed files with 185 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
Cargo.lock generated
View File

@ -17,6 +17,12 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"
[[package]]
name = "anyhow"
version = "1.0.71"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c7d0618f0e0b7e8ff11427422b64564d5fb0be1940354bfe2e0529b18a9d9b8"
[[package]]
name = "autocfg"
version = "1.1.0"
@ -1133,6 +1139,7 @@ checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
name = "screencap-bot"
version = "0.1.0"
dependencies = [
"anyhow",
"dotenvy",
"eggbug",
"env_logger",

1
Cargo.toml
View File

@ -6,6 +6,7 @@ authors = ["xenofem <xenofem@xeno.science>"]
license = "MIT"
[dependencies]
anyhow = "1.0.71"
dotenvy = "0.15.7"
eggbug = { git = "https://github.com/iliana/eggbug-rs.git", branch = "main" }
env_logger = "0.10"

3
README.md
View File

@ -19,7 +19,8 @@ directory:
- `SCREENCAP_BOT_SHOWS_FILE`: path of a YAML file specifying what shows to take screencaps from (default: `./shows.yaml`)
- `SCREENCAP_BOT_GLOBAL_TAGS`: tags to put on every post the bot makes, as a comma-separated list (eg `bot account,automated post,The Cohost Bot Feed`) (default: none)
- `SCREENCAP_BOT_POST_INTERVAL`: the interval between posts, in seconds (default: 6 hours)
- `SCREENCAP_BOT_POST_INTERVAL`: the interval between posts, in
seconds (default: 0, post a single screencap and then exit)
- `SCREENCAP_BOT_COHOST_EMAIL`: the email address the bot should use to log into cohost
- `SCREENCAP_BOT_COHOST_PASSWORD`: the password the bot should use to log into cohost
- `SCREENCAP_BOT_COHOST_PAGE`: the cohost page the bot should post from

2
src/config.rs
View File

@ -38,7 +38,7 @@ pub fn load() -> Config {
global_tags: get_var("GLOBAL_TAGS")
.map(|s| s.split(',').map(String::from).collect())
.unwrap_or_default(),
post_interval: Duration::from_secs(parse_var("POST_INTERVAL").unwrap_or(6 * 3600)),
post_interval: Duration::from_secs(parse_var("POST_INTERVAL").unwrap_or_default()),
cohost_email: expect_var("COHOST_EMAIL"),
cohost_password: expect_var("COHOST_PASSWORD"),
cohost_page: expect_var("COHOST_PAGE"),

214
src/main.rs
View File

@ -1,8 +1,11 @@
use std::time::Duration;
use anyhow::{anyhow, Context};
use config::Config;
use lazy_static::lazy_static;
use log::{debug, error, info};
use rand::{distributions::Standard, seq::IteratorRandom, Rng};
use shows::Shows;
mod config;
mod shows;
@ -15,7 +18,7 @@ lazy_static! {
}
#[tokio::main]
async fn main() {
async fn main() -> anyhow::Result<()> {
dotenvy::dotenv().ok();
env_logger::init();
ffmpeg_next::init().unwrap();
@ -25,114 +28,131 @@ async fn main() {
let conf = config::load();
info!("Loading shows from {}", conf.shows_file.display());
let shows = shows::load(conf.shows_file);
let shows = shows::load(&conf.shows_file).with_context(|| {
format!(
"Failed to load shows from file {}",
conf.shows_file.display()
)
})?;
if shows.is_empty() {
return Err(anyhow!("Shows file is empty!"));
}
info!("Logging into cohost as {}", conf.cohost_email);
let session = eggbug::Session::login(&conf.cohost_email, &conf.cohost_password)
.await
.expect("Failed to login to cohost");
.context("Failed to login to cohost")?;
loop {
let (title, show) = shows.iter().choose(&mut rng).expect("No shows found!");
let episodes = match show.episodes() {
Ok(eps) => eps,
Err(e) => {
error!("Failed to get episodes for {}: {}", title, e);
tokio::time::sleep(*RETRY_INTERVAL).await;
continue;
}
};
let (num, file) = episodes.iter().choose(&mut rng).unwrap();
let descriptor = format!(
"{}{}",
title,
match num {
EpisodeNumber::Standalone => String::new(),
EpisodeNumber::SingleSeason(n) => format!(" episode {}", n),
EpisodeNumber::MultiSeason(season, ep) =>
format!(" season {} episode {}", season, ep),
}
);
info!("Selected: {} - {}", descriptor, file.display());
let video_info = video::get_video_info(file, Some("eng")).unwrap();
debug!(
"Video duration: {}",
format_timestamp(video_info.duration_secs, None)
);
debug!(
"Subtitle stream index: {:?}",
video_info.subtitle_stream_index
);
let timestamp = video_info.duration_secs * rng.sample::<f64, _>(Standard);
let formatted_timestamp = format_timestamp(timestamp, Some(video_info.duration_secs));
info!("Taking screencap at {}", formatted_timestamp);
let img_data =
match video::take_screencap(file, timestamp, video_info.subtitle_stream_index).await {
Ok(data) => data,
Err(e) => {
error!("Failed to take screencap: {}", e);
tokio::time::sleep(*RETRY_INTERVAL).await;
continue;
}
};
let img_size = match imagesize::blob_size(&img_data) {
Ok(size) => size,
Err(e) => {
error!("Failed to get image size: {}", e);
tokio::time::sleep(*RETRY_INTERVAL).await;
continue;
}
};
let attachment = eggbug::Attachment::new(
img_data,
format!("{} @{}.png", descriptor, formatted_timestamp),
String::from("image/png"),
Some(img_size.width as u32),
Some(img_size.height as u32),
)
.with_alt_text(format!(
"Screencap of {} at {}",
descriptor, formatted_timestamp
));
let mut tags = show.tags.clone();
tags.extend_from_slice(&conf.global_tags);
match session
.create_post(
&conf.cohost_page,
&mut eggbug::Post {
content_warnings: vec![descriptor],
attachments: vec![attachment],
tags,
draft: false,
adult_content: false,
headline: String::new(),
markdown: String::new(),
metadata: None,
},
)
let result = post_random_screencap(&conf, &shows, &session, &mut rng)
.await
{
Ok(id) => info!("Created post {}", id),
Err(e) => {
error!("Failed to create post: {}", e);
tokio::time::sleep(*RETRY_INTERVAL).await;
continue;
}
.context("Failed to post a random screencap");
if conf.post_interval == Duration::ZERO {
return result;
}
tokio::time::sleep(conf.post_interval).await;
let delay = match result {
Ok(()) => conf.post_interval,
Err(e) => {
error!("{}", e);
*RETRY_INTERVAL
}
};
tokio::time::sleep(delay).await;
}
}
async fn post_random_screencap<R: Rng>(
conf: &Config,
shows: &Shows,
session: &eggbug::Session,
rng: &mut R,
) -> anyhow::Result<()> {
let (title, show) = shows.iter().choose(rng).unwrap();
let episodes = show.episodes().with_context(|| {
format!(
"Failed to get episode list for show {} with path {}",
title,
show.path.display()
)
})?;
let (num, file) = episodes.iter().choose(rng).unwrap();
let descriptor = format!(
"{}{}",
title,
match num {
EpisodeNumber::Standalone => String::new(),
EpisodeNumber::SingleSeason(n) => format!(" episode {}", n),
EpisodeNumber::MultiSeason(season, ep) => format!(" season {} episode {}", season, ep),
}
);
info!("Selected: {} - {}", descriptor, file.display());
let video_info = video::get_video_info(file, Some("eng")).with_context(|| {
format!(
"Failed to get duration and subtitle stream index for video file {}",
file.display()
)
})?;
debug!(
"Video duration: {}",
format_timestamp(video_info.duration_secs, None)
);
debug!(
"Subtitle stream index: {:?}",
video_info.subtitle_stream_index
);
let timestamp = video_info.duration_secs * rng.sample::<f64, _>(Standard);
let formatted_timestamp = format_timestamp(timestamp, Some(video_info.duration_secs));
info!("Taking screencap at {}", formatted_timestamp);
let img_data = video::take_screencap(file, timestamp, video_info.subtitle_stream_index)
.await
.context("Failed to take screencap")?;
let img_size = imagesize::blob_size(&img_data)
.context("Failed to get image size for screencap image data")?;
let attachment = eggbug::Attachment::new(
img_data,
format!("{} @{}.png", descriptor, formatted_timestamp),
String::from("image/png"),
Some(img_size.width as u32),
Some(img_size.height as u32),
)
.with_alt_text(format!(
"Screencap of {} at {}",
descriptor, formatted_timestamp
));
let mut tags = show.tags.clone();
tags.extend_from_slice(&conf.global_tags);
let chost_id = session
.create_post(
&conf.cohost_page,
&mut eggbug::Post {
content_warnings: vec![descriptor],
attachments: vec![attachment],
tags,
draft: false,
adult_content: false,
headline: String::new(),
markdown: String::new(),
metadata: None,
},
)
.await
.context("Failed to create chost")?;
info!("Created post {}", chost_id);
Ok(())
}
fn format_timestamp(timestamp: f64, total_duration: Option<f64>) -> String {
let total_duration = total_duration.unwrap_or(timestamp);
format!(

50
src/shows/mod.rs
View File

@ -1,5 +1,10 @@
use std::{collections::HashMap, fs, io::ErrorKind, path::PathBuf};
use std::{
collections::HashMap,
fs,
path::{Path, PathBuf},
};
use anyhow::{anyhow, Context};
use log::{debug, error};
use serde::Deserialize;
@ -12,6 +17,8 @@ pub struct Show {
pub tags: Vec<String>,
}
pub type Shows = HashMap<String, Show>;
#[derive(Debug, Eq, Hash, PartialEq)]
pub enum EpisodeNumber {
Standalone,
@ -21,15 +28,16 @@ pub enum EpisodeNumber {
type Episodes = HashMap<EpisodeNumber, PathBuf>;
pub fn load(shows_file: PathBuf) -> HashMap<String, Show> {
serde_yaml::from_reader(fs::File::open(shows_file).expect("Failed to open shows file"))
.expect("Failed to parse YAML from shows file")
pub fn load<P: AsRef<Path>>(shows_file: P) -> anyhow::Result<Shows> {
serde_yaml::from_reader(fs::File::open(shows_file).context("Failed to open shows file")?)
.context("Failed to parse YAML from shows file")
}
impl Show {
pub fn episodes(&self) -> std::io::Result<Episodes> {
pub fn episodes(&self) -> anyhow::Result<Episodes> {
let path = &self.path;
let metadata = fs::metadata(path)?;
let metadata =
fs::metadata(path).context("Failed to stat the show's path to determine file type")?;
if metadata.is_file() {
debug!("{} is a file, standalone", path.display());
Ok(HashMap::from([(
@ -38,10 +46,21 @@ impl Show {
)]))
} else if metadata.is_dir() {
debug!("{} is a directory, enumerating episodes", path.display());
let files: Vec<PathBuf> = fs::read_dir(path)?
let files: Vec<PathBuf> = fs::read_dir(path)
.context("Failed to get a directory listing for the show's path")?
.map(|entry| {
let entry = entry?;
if !entry.file_type()?.is_file() {
let entry =
entry.context("Failed to read a directory entry under the show's path")?;
if !entry
.file_type()
.with_context(|| {
format!(
"Failed to get file type for directory entry {}",
entry.path().display()
)
})?
.is_file()
{
debug!("Skipping {}, not a file", entry.path().display());
return Ok(None);
}
@ -55,16 +74,11 @@ impl Show {
Ok(Some(entry.path()))
})
.filter_map(|r| r.transpose())
.collect::<std::io::Result<Vec<PathBuf>>>()?;
enumeration::enumerate_episodes(files).ok_or(std::io::Error::new(
ErrorKind::InvalidData,
"No valid prefixes found",
))
.collect::<anyhow::Result<Vec<PathBuf>>>()?;
enumeration::enumerate_episodes(files)
.ok_or(anyhow!("Could not detect any episode numbering scheme"))
} else {
Err(std::io::Error::new(
ErrorKind::InvalidInput,
format!("Invalid file type for {}", path.display()),
))
Err(anyhow!("The show's path is not a file or a directory"))
}
}
}

40
src/video.rs
View File

@ -1,3 +1,4 @@
use anyhow::{anyhow, Context};
use ffmpeg_next::{
format::{input, stream::Disposition},
media::Type,
@ -7,7 +8,7 @@ use regex::Regex;
use tempfile::tempdir;
use tokio::{fs, process::Command};
use std::{io::ErrorKind, path::Path};
use std::path::Path;
lazy_static! {
static ref SUBTITLE_FORBID_REGEX: Regex = Regex::new("(?i)sign|song").unwrap();
@ -22,8 +23,8 @@ pub struct VideoInfo {
pub fn get_video_info<P: AsRef<Path>>(
source: &P,
subtitle_lang: Option<&str>,
) -> Result<VideoInfo, ffmpeg_next::Error> {
let ctx = input(source)?;
) -> anyhow::Result<VideoInfo> {
let ctx = input(source).context("Failed to load video file")?;
let duration_secs = ctx.duration() as f64 / f64::from(ffmpeg_next::ffi::AV_TIME_BASE);
@ -63,18 +64,21 @@ pub async fn take_screencap<P: AsRef<Path>>(
source: &P,
timestamp_secs: f64,
subtitle_stream_index: Option<usize>,
) -> std::io::Result<Vec<u8>> {
) -> anyhow::Result<Vec<u8>> {
let ext = source.as_ref().extension().and_then(|s| s.to_str());
if ext != Some("mkv") && ext != Some("mp4") {
return Err(std::io::Error::new(
ErrorKind::Other,
"unexpected file extension",
return Err(anyhow!(
"Video file {} had unexpected file extension",
source.as_ref().display()
));
}
let tmp_dir = tempdir()?;
let tmp_dir = tempdir()
.context("Failed to create temporary directory for ffmpeg input and output files")?;
let link_path = tmp_dir.path().join(format!("in.{}", ext.unwrap()));
fs::symlink(source, &link_path).await?;
fs::symlink(source, &link_path)
.await
.context("Failed to create symlink for video file")?;
let dest_path = tmp_dir.path().join("out.png");
let mut cmd = Command::new("ffmpeg");
@ -98,12 +102,18 @@ pub async fn take_screencap<P: AsRef<Path>>(
.arg("-y")
.arg(&dest_path);
if !cmd.status().await?.success() {
return Err(std::io::Error::new(
ErrorKind::Other,
"ffmpeg command failed",
));
let status = cmd
.status()
.await
.with_context(|| format!("Error running ffmpeg child process {:?}", cmd))?;
if !status.success() {
match status.code() {
Some(code) => return Err(anyhow!("ffmpeg exited with status code {code}")),
None => return Err(anyhow!("ffmpeg terminated by signal")),
}
}
fs::read(&dest_path).await
fs::read(&dest_path)
.await
.with_context(|| format!("Failed to read ffmpeg output file {}", dest_path.display()))
}