1 files changed, 49 insertions, 19 deletions

diff --git a/src/main.rs b/src/main.rs
index 96455ac..f427d6c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,6 @@
 use anyhow::Result;
 use async_trait::async_trait;
-// use clap::{Parser, ValueEnum};
+use clap::Parser;
 use env_logger;
 use image::GenericImageView;
 use image::ImageReader;
@@ -13,24 +13,43 @@ use rfb::rfb::{
 use rfb::server::{Server, VncServer, VncServerConfig, VncServerData};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
-const WIDTH: usize = 1053;
-const HEIGHT: usize = 1053;
-
 #[derive(Clone)]
 struct AcheServer {
     rgb_order: (u8, u8, u8),
     big_endian: bool,
+    screen: String,
+}
+
+#[derive(Parser)]
+#[command(version, about, long_about = None)]
+struct Args {
+    /// The port to listen to
+    #[arg(short, long, default_value_t = 5900)]
+    port: u16,
+
+    /// The picture to show
+    #[arg(short, long)]
+    screen: String,
+
+    /// The server advertised name
+    #[arg(short, long, default_value_t = String::from("ache-hacked-server"))]
+    name_string: String, // 'name-string' is the server name in the RFC
 }
 
 #[tokio::main]
 async fn main() -> Result<()> {
     env_logger::init();
+
     /*
-        TODO: Should add a CLI option to configure the port to listen to.
+        TODO: Should log keys press events.
 
-        TODO: Should add a CLI option to configure the image to display.
+        TODO: Only read the image once.
+
+        TODO: Improve String management.
     */
 
+    let args = Args::parse();
+
     let pf = PixelFormat::new_colorformat(
         rgb_888::BITS_PER_PIXEL,
         rgb_888::DEPTH,
@@ -45,19 +64,25 @@ async fn main() -> Result<()> {
     info!("Starting server: pixel format; {:#?}", pf);
 
     let config = VncServerConfig {
-        addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 5901),
+        addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), args.port),
         version: ProtoVersion::Rfb38,
         sec_types: SecurityTypes(vec![SecurityType::None, SecurityType::VncAuthentication]),
-        name: "rfb-ache-server".to_string(),
+        name: args.name_string,
     };
+    let dim = ImageReader::open(args.screen.clone())
+        .unwrap()
+        .into_dimensions()
+        .unwrap();
+
     let data = VncServerData {
-        width: WIDTH as u16,
-        height: HEIGHT as u16,
+        width: dim.0 as u16,
+        height: dim.1 as u16,
         input_pixel_format: pf.clone(),
     };
     let server = AcheServer {
         rgb_order: (0, 1, 2),
         big_endian: false,
+        screen: args.screen.clone(),
     };
     let s = VncServer::new(server, config, data);
     s.start().await?;
@@ -80,11 +105,17 @@ fn order_to_index(order: u8, big_endian: bool) -> u8 {
     }
 }
 
-fn generate_image(name: &str, big_endian: bool, rgb_order: (u8, u8, u8)) -> Vec<u8> {
-    const LEN: usize = WIDTH * HEIGHT * rgb_888::BYTES_PER_PIXEL;
-    let mut pixels = vec![0xffu8; LEN];
-
+fn generate_image(name: &str, big_endian: bool, rgb_order: (u8, u8, u8)) -> (Vec<u8>, u16, u16) {
     let img = ImageReader::open(name).unwrap().decode().unwrap();
+    let (width, height) = img.dimensions();
+    let width = width as usize;
+    let height = height as usize;
+
+    let width_u16 = u16::try_from(width).unwrap();
+    let height_u16 = u16::try_from(height).unwrap();
+
+    let len: usize = width * height * rgb_888::BYTES_PER_PIXEL;
+    let mut pixels = vec![0x11u8; len];
 
     let (r, g, b) = rgb_order;
     let r_idx = order_to_index(r, big_endian) as usize;
@@ -97,7 +128,7 @@ fn generate_image(name: &str, big_endian: bool, rgb_order: (u8, u8, u8)) -> Vec<
         let ux = x as usize;
         let uy = y as usize;
 
-        let y_offset = WIDTH * rgb_888::BYTES_PER_PIXEL;
+        let y_offset = width * rgb_888::BYTES_PER_PIXEL;
         let x_offset = ux * rgb_888::BYTES_PER_PIXEL;
 
         pixels[uy * y_offset + x_offset + r_idx] = pixel[0];
@@ -106,15 +137,14 @@ fn generate_image(name: &str, big_endian: bool, rgb_order: (u8, u8, u8)) -> Vec<
         pixels[uy * y_offset + x_offset + x_idx] = pixel[3];
     }
 
-    pixels
+    (pixels, width_u16, height_u16)
 }
 
 #[async_trait]
 impl Server for AcheServer {
     async fn get_framebuffer_update(&self) -> FramebufferUpdate {
-        let pixels_width = 1053;
-        let pixels_height = 1053;
-        let pixels = generate_image("screen-rfp.png", self.big_endian, self.rgb_order);
+        let (pixels, pixels_width, pixels_height) =
+            generate_image(&self.screen, self.big_endian, self.rgb_order);
         let r = Rectangle::new(
             0,
             0,