Proof-of-concept: takes a random screencap, doesn't post it yet

.gitignore

@@ -0,0 +1 @@
+/target

Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "screencap-bot"
+version = "0.1.0"
+edition = "2021"
+authors = ["xenofem <xenofem@xeno.science>"]
+license = "MIT"
+
+[dependencies]
+dotenvy = "0.15.7"
+eggbug = "0.1.3"
+env_logger = "0.10"
+ffmpeg-next = "6.0.0"
+lazy_static = "1.4.0"
+log = "0.4.19"
+rand = "0.8"
+regex = "1.8.4"
+serde = "1"
+serde_yaml = "0.9.22"
+tempfile = "3.6.0"
+tokio = { version = "1.28.2", features = ["full"] }

flake.lock

@@ -0,0 +1,24 @@
+{
+  "nodes": {
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1686226982,
+        "narHash": "sha256-nLuiPoeiVfqqzeq9rmXxpybh77VS37dsY/k8N2LoxVg=",
+        "path": "/nix/store/xwqh0l2x34wfn4yxpgx8f4qfzhm56rnv-source",
+        "rev": "a64b73e07d4aa65cfcbda29ecf78eaf9e72e44bd",
+        "type": "path"
+      },
+      "original": {
+        "id": "nixpkgs",
+        "type": "indirect"
+      }
+    },
+    "root": {
+      "inputs": {
+        "nixpkgs": "nixpkgs"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,15 @@
+{
+  description = "Screencap bot for Cohost";
+
+  outputs = { self, nixpkgs }: let
+    pkgs = import nixpkgs {
+      system = "x86_64-linux";
+    };
+  in {
+    devShells.x86_64-linux.default = pkgs.mkShell {
+      buildInputs = with pkgs; [ ffmpeg openssl ];
+      nativeBuildInputs = with pkgs; [ rustc cargo pkgconfig clang ];
+      LIBCLANG_PATH = "${pkgs.llvmPackages.libclang.lib}/lib";
+    };
+  };
+}

src/config.rs

@@ -0,0 +1,17 @@
+use std::{collections::HashMap, env, fs, path::PathBuf};
+
+pub type Shows = HashMap<String, PathBuf>;
+
+pub struct Config {
+    pub shows: Shows,
+}
+
+pub fn load() -> Config {
+    let shows_file = env::var("SCREENCAP_BOT_SHOWS_FILE").unwrap_or(String::from("./shows.yaml"));
+    Config {
+        shows: serde_yaml::from_reader(
+            fs::File::open(shows_file).expect("Failed to open shows file"),
+        )
+        .expect("Failed to parse YAML from shows file"),
+    }
+}

src/enumeration.rs

@@ -0,0 +1,122 @@
+use std::{
+    collections::{BTreeSet, HashMap},
+    fs,
+    io::ErrorKind,
+    path::{Path, PathBuf},
+};
+
+use lazy_static::lazy_static;
+use log::{debug, error, trace, warn};
+use regex::Regex;
+
+use crate::config::Shows;
+
+lazy_static! {
+    static ref NUMBER_REGEX: Regex = Regex::new("[0-9]+").unwrap();
+}
+
+/// [episode number, or None for a standalone movie] -> [video file path]
+type Enumeration = HashMap<Option<u32>, PathBuf>;
+
+/// [show name] -> [enumeration]
+type EnumeratedShows = HashMap<String, Enumeration>;
+
+pub fn enumerate_shows(shows: Shows) -> EnumeratedShows {
+    shows
+        .into_iter()
+        .filter_map(|(name, path)| {
+            debug!("Enumerating {} from {}", name, path.display());
+            enumerate_show(path)
+                .map_err(|e| {
+                    error!("Error processing {}: {}", name, e);
+                })
+                .ok()
+                .map(|en| (name, en))
+        })
+        .collect()
+}
+
+fn enumerate_show(path: PathBuf) -> std::io::Result<Enumeration> {
+    let metadata = fs::metadata(&path)?;
+    if metadata.is_file() {
+        debug!("{} is a file, standalone", path.display());
+        Ok(HashMap::from([(None, path)]))
+    } else if metadata.is_dir() {
+        debug!("{} is a directory, enumerating episodes", path.display());
+        enumerate_dir(path)
+    } else {
+        Err(std::io::Error::new(
+            ErrorKind::InvalidInput,
+            format!("Invalid file type for {}", path.display()),
+        ))
+    }
+}
+
+fn enumerate_dir(path: PathBuf) -> std::io::Result<Enumeration> {
+    let files: Vec<(PathBuf, String)> = fs::read_dir(&path)?
+        .map(|entry| {
+            let entry = entry?;
+            if !entry.file_type()?.is_file() {
+                debug!("Skipping {}, not a file", entry.path().display());
+                return Ok(None);
+            }
+            Ok(Some((
+                entry.path(),
+                entry.file_name().into_string().map_err(|o| {
+                    std::io::Error::new(ErrorKind::InvalidData, format!("Invalid filename {:?}", o))
+                })?,
+            )))
+        })
+        .filter_map(|r| r.transpose())
+        .collect::<std::io::Result<Vec<(PathBuf, String)>>>()?;
+
+    // For each string prefix which is followed by a number in one or
+    // more filenames, record what the number is for each filename
+    // with that prefix, sorted by number.
+    //
+    // We use a BTreeSet of pairs rather than a BTreeMap so that
+    // duplicates are preserved.
+    let mut prefix_enumerations: HashMap<&str, BTreeSet<(u32, &Path)>> = HashMap::new();
+    for (path, filename) in files.iter() {
+        for m in NUMBER_REGEX.find_iter(filename) {
+            let num = m.as_str().parse();
+            if let Err(e) = num {
+                warn!("Failed to parse candidate number: {}", e);
+                continue;
+            }
+            let num = num.unwrap();
+            let prefix = filename.get(..m.start()).unwrap();
+            trace!("{}: candidate prefix {}, number {}", filename, prefix, num);
+            prefix_enumerations
+                .entry(prefix)
+                .or_default()
+                .insert((num, path));
+        }
+    }
+
+    // Use the prefix enumeration that produces the most consecutive numbers
+    let best_enumeration = prefix_enumerations
+        .into_values()
+        .max_by_key(|set| {
+            set.iter()
+                .zip(set.iter().skip(1))
+                .filter(|((cur, _), (next, _))| next - cur == 1)
+                .count()
+        })
+        .ok_or(std::io::Error::new(
+            ErrorKind::InvalidData,
+            "No valid prefixes found",
+        ))?;
+
+    let mut result: Enumeration = HashMap::new();
+    for (num, path) in best_enumeration.into_iter() {
+        debug!("Episode {}: {}", num, path.display());
+        if let Some(dup) = result.insert(Some(num), path.to_path_buf()) {
+            warn!(
+                "Duplicate episode number, discarding file {}",
+                dup.display()
+            );
+        }
+    }
+    Ok(result)
+}

src/main.rs

@@ -0,0 +1,80 @@
+use std::path::Path;
+
+use enumeration::enumerate_shows;
+use log::{debug, error, info};
+use rand::{distributions::Standard, seq::IteratorRandom, Rng};
+
+mod config;
+mod enumeration;
+mod video;
+
+#[tokio::main]
+async fn main() -> std::io::Result<()> {
+    env_logger::init();
+    ffmpeg_next::init().unwrap();
+    dotenvy::dotenv().ok();
+
+    let mut rng = rand::thread_rng();
+
+    let conf = config::load();
+    let enumerated_shows = enumerate_shows(conf.shows);
+
+    loop {
+        let (show, episodes) = enumerated_shows
+            .iter()
+            .choose(&mut rng)
+            .expect("No shows found!");
+        let (num, file) = episodes.iter().choose(&mut rng).unwrap();
+
+        let descriptor = format!(
+            "{}{}",
+            show,
+            if let Some(n) = num {
+                format!(" episode {}", n)
+            } else {
+                String::new()
+            }
+        );
+
+        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);
+        if let Err(e) = video::take_screencap(
+            file,
+            timestamp,
+            video_info.subtitle_stream_index,
+            &Path::new("./out.png"),
+        ) {
+            error!("Failed to take screencap: {}", e);
+        }
+
+        break;
+    }
+    Ok(())
+}
+
+fn format_timestamp(timestamp: f64, total_duration: Option<f64>) -> String {
+    let total_duration = total_duration.unwrap_or(timestamp);
+    format!(
+        "{}{:02}:{:05.2}",
+        if total_duration >= 3600.0 {
+            format!("{}:", (timestamp / 3600.0) as u32)
+        } else {
+            String::new()
+        },
+        ((timestamp % 3600.0) / 60.0) as u32,
+        timestamp % 60.0
+    )
+}

src/video.rs

@@ -0,0 +1,107 @@
+use ffmpeg_next::{
+    format::{input, stream::Disposition},
+    media::Type,
+};
+use lazy_static::lazy_static;
+use regex::Regex;
+use tempfile::tempdir;
+
+use std::{ffi::OsStr, io::ErrorKind, os::unix::fs, path::Path, process::Command};
+
+lazy_static! {
+    static ref SUBTITLE_FORBID_REGEX: Regex = Regex::new("(?i)sign|song").unwrap();
+}
+
+pub struct VideoInfo {
+    pub duration_secs: f64,
+    // The index among the subtitle streams, not among the streams in general
+    pub subtitle_stream_index: Option<usize>,
+}
+
+pub fn get_video_info<P: AsRef<Path>>(
+    source: &P,
+    subtitle_lang: Option<&str>,
+) -> Result<VideoInfo, ffmpeg_next::Error> {
+    let ctx = input(source)?;
+
+    let duration_secs = ctx.duration() as f64 / f64::from(ffmpeg_next::ffi::AV_TIME_BASE);
+
+    let subtitle_stream_index = subtitle_lang
+        .map(|lang| {
+            ctx.streams()
+                .filter(|stream| {
+                    ffmpeg_next::codec::context::Context::from_parameters(stream.parameters())
+                        .map(|c| c.medium())
+                        == Ok(Type::Subtitle)
+                })
+                .enumerate()
+                .filter(|(_, stream)| {
+                    let metadata = stream.metadata();
+                    if metadata.get("language") != Some(lang) {
+                        return false;
+                    }
+                    if metadata
+                        .get("title")
+                        .map(|t| SUBTITLE_FORBID_REGEX.is_match(t))
+                        == Some(true)
+                    {
+                        return false;
+                    }
+                    true
+                })
+                .min_by_key(|(_, stream)| stream.disposition().contains(Disposition::FORCED))
+                .map(|(idx, _)| idx)
+        })
+        .flatten();
+
+    Ok(VideoInfo {
+        duration_secs,
+        subtitle_stream_index,
+    })
+}
+
+pub fn take_screencap<P: AsRef<Path>, Q: AsRef<OsStr>>(
+    source: &P,
+    timestamp_secs: f64,
+    subtitle_stream_index: Option<usize>,
+    dest: &Q,
+) -> std::io::Result<()> {
+    let ext = source.as_ref().extension().map(|s| s.to_str()).flatten();
+    if ext != Some("mkv") && ext != Some("mp4") {
+        return Err(std::io::Error::new(
+            ErrorKind::Other,
+            "unexpected file extension",
+        ));
+    }
+    let tmp_dir = tempdir()?;
+    let link_path = tmp_dir.path().join(format!("input.{}", ext.unwrap()));
+    fs::symlink(source, &link_path)?;
+    let mut cmd = Command::new("ffmpeg");
+
+    cmd.arg("-ss")
+        .arg(format!("{:.2}", timestamp_secs))
+        .arg("-copyts")
+        .arg("-i")
+        .arg(&link_path);
+
+    if let Some(idx) = subtitle_stream_index {
+        cmd.arg("-filter_complex").arg(format!(
+            "[0:v]subtitles={}:si={}",
+            link_path.to_string_lossy(),
+            idx
+        ));
+    }
+
+    cmd.args(["-vframes", "1"])
+        .args(["-loglevel", "quiet"])
+        .arg("-y")
+        .arg(&dest);
+
+    if !cmd.status()?.success() {
+        return Err(std::io::Error::new(
+            ErrorKind::Other,
+            "ffmpeg command failed",
+        ));
+    }
+    Ok(())
+}