mirror of
https://github.com/rustdesk/rustdesk.git
synced 2024-12-15 03:59:04 +08:00
157 lines
4.1 KiB
Rust
157 lines
4.1 KiB
Rust
extern crate docopt;
|
|
extern crate quest;
|
|
extern crate repng;
|
|
extern crate scrap;
|
|
extern crate serde;
|
|
extern crate webm;
|
|
|
|
use std::fs::{File, OpenOptions};
|
|
use std::path::PathBuf;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use std::{io, thread};
|
|
|
|
use docopt::Docopt;
|
|
use scrap::codec::{EncoderApi, EncoderCfg};
|
|
use webm::mux;
|
|
use webm::mux::Track;
|
|
|
|
use scrap::vpxcodec as vpx_encode;
|
|
use scrap::{TraitCapturer, Capturer, Display, STRIDE_ALIGN};
|
|
|
|
const USAGE: &'static str = "
|
|
Simple WebM screen capture.
|
|
|
|
Usage:
|
|
record-screen <path> [--time=<s>] [--fps=<fps>] [--bv=<kbps>] [--ba=<kbps>] [--codec CODEC]
|
|
record-screen (-h | --help)
|
|
|
|
Options:
|
|
-h --help Show this screen.
|
|
--time=<s> Recording duration in seconds.
|
|
--fps=<fps> Frames per second [default: 30].
|
|
--bv=<kbps> Video bitrate in kilobits per second [default: 5000].
|
|
--ba=<kbps> Audio bitrate in kilobits per second [default: 96].
|
|
--codec CODEC Configure the codec used. [default: vp9]
|
|
Valid values: vp8, vp9.
|
|
";
|
|
|
|
#[derive(Debug, serde::Deserialize)]
|
|
struct Args {
|
|
arg_path: PathBuf,
|
|
flag_codec: Codec,
|
|
flag_time: Option<u64>,
|
|
flag_fps: u64,
|
|
flag_bv: u32,
|
|
}
|
|
|
|
#[derive(Debug, serde::Deserialize)]
|
|
enum Codec {
|
|
Vp8,
|
|
Vp9,
|
|
}
|
|
|
|
fn main() -> io::Result<()> {
|
|
let args: Args = Docopt::new(USAGE)
|
|
.and_then(|d| d.deserialize())
|
|
.unwrap_or_else(|e| e.exit());
|
|
|
|
let duration = args.flag_time.map(Duration::from_secs);
|
|
|
|
let d = Display::primary().unwrap();
|
|
let (width, height) = (d.width() as u32, d.height() as u32);
|
|
|
|
// Setup the multiplexer.
|
|
|
|
let out = match {
|
|
OpenOptions::new()
|
|
.write(true)
|
|
.create_new(true)
|
|
.open(&args.arg_path)
|
|
} {
|
|
Ok(file) => file,
|
|
Err(ref e) if e.kind() == io::ErrorKind::AlreadyExists => {
|
|
if loop {
|
|
quest::ask("Overwrite the existing file? [y/N] ");
|
|
if let Some(b) = quest::yesno(false)? {
|
|
break b;
|
|
}
|
|
} {
|
|
File::create(&args.arg_path)?
|
|
} else {
|
|
return Ok(());
|
|
}
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
|
|
let mut webm =
|
|
mux::Segment::new(mux::Writer::new(out)).expect("Could not initialize the multiplexer.");
|
|
|
|
let (vpx_codec, mux_codec) = match args.flag_codec {
|
|
Codec::Vp8 => (vpx_encode::VpxVideoCodecId::VP8, mux::VideoCodecId::VP8),
|
|
Codec::Vp9 => (vpx_encode::VpxVideoCodecId::VP9, mux::VideoCodecId::VP9),
|
|
};
|
|
|
|
let mut vt = webm.add_video_track(width, height, None, mux_codec);
|
|
|
|
// Setup the encoder.
|
|
|
|
let mut vpx = vpx_encode::VpxEncoder::new(EncoderCfg::VPX(vpx_encode::VpxEncoderConfig {
|
|
width,
|
|
height,
|
|
timebase: [1, 1000],
|
|
bitrate: args.flag_bv,
|
|
codec: vpx_codec,
|
|
num_threads: 0,
|
|
}))
|
|
.unwrap();
|
|
|
|
// Start recording.
|
|
|
|
let start = Instant::now();
|
|
let stop = Arc::new(AtomicBool::new(false));
|
|
|
|
thread::spawn({
|
|
let stop = stop.clone();
|
|
move || {
|
|
let _ = quest::ask("Recording! Press ⏎ to stop.");
|
|
let _ = quest::text();
|
|
stop.store(true, Ordering::Release);
|
|
}
|
|
});
|
|
|
|
let spf = Duration::from_nanos(1_000_000_000 / args.flag_fps);
|
|
|
|
// Capturer object is expensive, avoiding to create it frequently.
|
|
let mut c = Capturer::new(d, true).unwrap();
|
|
while !stop.load(Ordering::Acquire) {
|
|
let now = Instant::now();
|
|
let time = now - start;
|
|
|
|
if Some(true) == duration.map(|d| time > d) {
|
|
break;
|
|
}
|
|
|
|
if let Ok(frame) = c.frame(Duration::from_millis(0)) {
|
|
let ms = time.as_secs() * 1000 + time.subsec_millis() as u64;
|
|
|
|
for frame in vpx.encode(ms as i64, &frame, STRIDE_ALIGN).unwrap() {
|
|
vt.add_frame(frame.data, frame.pts as u64 * 1_000_000, frame.key);
|
|
}
|
|
}
|
|
|
|
let dt = now.elapsed();
|
|
if dt < spf {
|
|
thread::sleep(spf - dt);
|
|
}
|
|
}
|
|
|
|
// End things.
|
|
|
|
let _ = webm.finalize(None);
|
|
|
|
Ok(())
|
|
}
|