1 files changed, 17 insertions, 6 deletions

diff --git a/src/main.rs b/src/main.rs
index b3635ee..cf5537d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,9 +11,17 @@ static PROFILE_ROOT: &str = "/tmp/geckodriver-dashboard";
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
+    /// Input HTML file.
+    #[arg(short, long)]
+    input_path: String,
+
     /// Output path for png image.
     #[arg(short, long)]
-    path: String,
+    output_path: String,
+
+    /// CSS locator to generate screenshot for.
+    #[arg(default_value_t = String::from("#app"))]
+    element_locator: String,
 }
 
 // let's set up the sequence of steps we want the browser to take
@@ -21,11 +29,14 @@ struct Args {
 async fn main() -> Result<(), Box<dyn Error>> {
     pretty_env_logger::init();
     let args = Args::parse();
-    let display = env::var("DISPLAY")?;
+    let display = ":5"; // env::var("DISPLAY")?;
+
+    let _xvfb = Command::new("Xvfb")
+        .arg(display).spawn()?;
 
     info!("Current directory: {}", env::current_dir()?.display());
     fs::create_dir(PROFILE_ROOT).ok();
-    let mut geckodriver = Command::new("geckodriver")
+    let _geckodriver = Command::new("geckodriver")
         .env("MOZ_REMOTE_SETTINGS_DEVTOOLS", "1")
         .env("DISPLAY", display)
         .arg("--profile-root").arg("/tmp/geckodriver-dashboard").spawn()?;
@@ -33,14 +44,14 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let c = ClientBuilder::native().connect("http://localhost:4444").await.expect("failed to connect to WebDriver");
 
     loop {
-        let path = fs::canonicalize("index.html")?;
+        let path = fs::canonicalize(&args.input_path)?;
         info!("Opening {}", path.display());
         c.goto(&format!("file://{}", path.display())).await?;
 
-        let e = c.find(Locator::Css("#app")).await?;
+        let e = c.find(Locator::Css(&args.element_locator)).await?;
         let img =  e.screenshot().await?;
 
-        let output_path = Path::new(&args.path);
+        let output_path = Path::new(&args.output_path);
         let tmp_path = output_path.with_extension("png.tmp");
         fs::write(&tmp_path, img)?;