Refactor code writer and fix minor things

Also added the ability to specify which files should be generated based on their file type.

E.g. `cs2-dumper.exe -f hpp,json`

src/output/mod.rs

@@ -1,4 +1,4 @@
-use std::fmt::Write;
+use std::fmt::{self, Write};
 use std::fs;
 use std::path::Path;
 
@@ -6,7 +6,7 @@ use chrono::{DateTime, Utc};
 
 use memflow::prelude::v1::*;
 
-use serde::{Deserialize, Serialize};
+use serde::Serialize;
 use serde_json::json;
 
 use formatter::Formatter;
@@ -30,216 +30,160 @@ enum Item<'a> {
 }
 
 impl<'a> Item<'a> {
-    fn generate(&self, results: &Results, indent_size: usize, file_ext: &str) -> Result<String> {
+    fn write(&self, fmt: &mut Formatter<'a>, file_ext: &str) -> fmt::Result {
         match file_ext {
-            "cs" => self.to_cs(results, indent_size),
-            "hpp" => self.to_hpp(results, indent_size),
-            "json" => self.to_json(results, indent_size),
-            "rs" => self.to_rs(results, indent_size),
-            _ => unreachable!(),
+            "cs" => self.write_cs(fmt),
+            "hpp" => self.write_hpp(fmt),
+            "json" => self.write_json(fmt),
+            "rs" => self.write_rs(fmt),
+            _ => unimplemented!(),
         }
     }
 }
 
-trait CodeGen {
-    fn to_cs(&self, results: &Results, indent_size: usize) -> Result<String>;
+trait CodeWriter {
+    fn write_cs(&self, fmt: &mut Formatter<'_>) -> fmt::Result;
 
-    fn to_hpp(&self, results: &Results, indent_size: usize) -> Result<String>;
+    fn write_hpp(&self, fmt: &mut Formatter<'_>) -> fmt::Result;
 
-    fn to_json(&self, results: &Results, indent_size: usize) -> Result<String>;
+    fn write_json(&self, fmt: &mut Formatter<'_>) -> fmt::Result;
 
-    fn to_rs(&self, results: &Results, indent_size: usize) -> Result<String>;
-
-    /// Replaces non-alphanumeric characters in a string with underscores.
-    #[inline]
-    fn slugify(input: &str) -> String {
-        input.replace(|c: char| !c.is_alphanumeric(), "_")
-    }
-
-    fn write_content<F>(&self, results: &Results, indent_size: usize, f: F) -> Result<String>
-    where
-        F: FnOnce(&mut Formatter<'_>) -> Result<()>,
-    {
-        let mut buf = String::new();
-        let mut fmt = Formatter::new(&mut buf, indent_size);
-
-        results.write_banner(&mut fmt)?;
-
-        f(&mut fmt)?;
-
-        Ok(buf)
-    }
+    fn write_rs(&self, fmt: &mut Formatter<'_>) -> fmt::Result;
 }
 
-impl<'a> CodeGen for Item<'a> {
-    fn to_cs(&self, results: &Results, indent_size: usize) -> Result<String> {
+impl<'a> CodeWriter for Item<'a> {
+    fn write_cs(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         match self {
-            Item::Buttons(buttons) => buttons.to_cs(results, indent_size),
-            Item::Interfaces(ifaces) => ifaces.to_cs(results, indent_size),
-            Item::Offsets(offsets) => offsets.to_cs(results, indent_size),
-            Item::Schemas(schemas) => schemas.to_cs(results, indent_size),
+            Item::Buttons(buttons) => buttons.write_cs(fmt),
+            Item::Interfaces(ifaces) => ifaces.write_cs(fmt),
+            Item::Offsets(offsets) => offsets.write_cs(fmt),
+            Item::Schemas(schemas) => schemas.write_cs(fmt),
         }
     }
 
-    fn to_hpp(&self, results: &Results, indent_size: usize) -> Result<String> {
+    fn write_hpp(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         match self {
-            Item::Buttons(buttons) => buttons.to_hpp(results, indent_size),
-            Item::Interfaces(ifaces) => ifaces.to_hpp(results, indent_size),
-            Item::Offsets(offsets) => offsets.to_hpp(results, indent_size),
-            Item::Schemas(schemas) => schemas.to_hpp(results, indent_size),
+            Item::Buttons(buttons) => buttons.write_hpp(fmt),
+            Item::Interfaces(ifaces) => ifaces.write_hpp(fmt),
+            Item::Offsets(offsets) => offsets.write_hpp(fmt),
+            Item::Schemas(schemas) => schemas.write_hpp(fmt),
         }
     }
 
-    fn to_json(&self, results: &Results, indent_size: usize) -> Result<String> {
+    fn write_json(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         match self {
-            Item::Buttons(buttons) => buttons.to_json(results, indent_size),
-            Item::Interfaces(ifaces) => ifaces.to_json(results, indent_size),
-            Item::Offsets(offsets) => offsets.to_json(results, indent_size),
-            Item::Schemas(schemas) => schemas.to_json(results, indent_size),
+            Item::Buttons(buttons) => buttons.write_json(fmt),
+            Item::Interfaces(ifaces) => ifaces.write_json(fmt),
+            Item::Offsets(offsets) => offsets.write_json(fmt),
+            Item::Schemas(schemas) => schemas.write_json(fmt),
         }
     }
 
-    fn to_rs(&self, results: &Results, indent_size: usize) -> Result<String> {
+    fn write_rs(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
         match self {
-            Item::Buttons(buttons) => buttons.to_rs(results, indent_size),
-            Item::Interfaces(ifaces) => ifaces.to_rs(results, indent_size),
-            Item::Offsets(offsets) => offsets.to_rs(results, indent_size),
-            Item::Schemas(schemas) => schemas.to_rs(results, indent_size),
+            Item::Buttons(buttons) => buttons.write_rs(fmt),
+            Item::Interfaces(ifaces) => ifaces.write_rs(fmt),
+            Item::Offsets(offsets) => offsets.write_rs(fmt),
+            Item::Schemas(schemas) => schemas.write_rs(fmt),
         }
     }
 }
 
-#[derive(Deserialize, Serialize)]
-pub struct Results {
-    /// Timestamp of the dump.
-    pub timestamp: DateTime<Utc>,
-
-    /// List of buttons to dump.
-    pub buttons: Vec<Button>,
-
-    /// Map of interfaces to dump.
-    pub interfaces: InterfaceMap,
-
-    /// Map of offsets to dump.
-    pub offsets: OffsetMap,
-
-    /// Map of schema classes/enums to dump.
-    pub schemas: SchemaMap,
+pub struct Output<'a> {
+    file_types: &'a Vec<String>,
+    indent_size: usize,
+    out_dir: &'a Path,
+    result: &'a AnalysisResult,
+    timestamp: DateTime<Utc>,
 }
 
-impl Results {
+impl<'a> Output<'a> {
     pub fn new(
-        buttons: Vec<Button>,
-        interfaces: InterfaceMap,
-        offsets: OffsetMap,
-        schemas: SchemaMap,
-    ) -> Self {
-        Self {
-            timestamp: Utc::now(),
-            buttons,
-            interfaces,
-            offsets,
-            schemas,
-        }
-    }
-
-    pub fn dump_all<P: AsRef<Path>>(
-        &self,
-        process: &mut IntoProcessInstanceArcBox<'_>,
-        out_dir: P,
+        file_types: &'a Vec<String>,
         indent_size: usize,
-    ) -> Result<()> {
-        // TODO: Make this user-configurable.
-        const FILE_EXTS: &[&str] = &["cs", "hpp", "json", "rs"];
-
+        out_dir: &'a Path,
+        result: &'a AnalysisResult,
+    ) -> Result<Self> {
         fs::create_dir_all(&out_dir)?;
 
+        Ok(Self {
+            file_types,
+            indent_size,
+            out_dir,
+            result,
+            timestamp: Utc::now(),
+        })
+    }
+
+    pub fn dump_all(&self, process: &mut IntoProcessInstanceArcBox<'_>) -> Result<()> {
         let items = [
-            ("buttons", Item::Buttons(&self.buttons)),
-            ("interfaces", Item::Interfaces(&self.interfaces)),
-            ("offsets", Item::Offsets(&self.offsets)),
+            ("buttons", Item::Buttons(&self.result.buttons)),
+            ("interfaces", Item::Interfaces(&self.result.interfaces)),
+            ("offsets", Item::Offsets(&self.result.offsets)),
         ];
 
         for (file_name, item) in &items {
-            self.dump_item(item, &out_dir, indent_size, FILE_EXTS, file_name)?;
+            self.dump_item(file_name, item)?;
         }
 
-        self.dump_info(process, &out_dir)?;
-        self.dump_schemas(&out_dir, indent_size, FILE_EXTS)?;
+        self.dump_schemas()?;
+        self.dump_info(process)?;
 
         Ok(())
     }
 
-    fn dump_file<P: AsRef<Path>>(
-        &self,
-        out_dir: P,
-        file_name: &str,
-        file_ext: &str,
-        content: &str,
-    ) -> Result<()> {
-        let file_path = out_dir.as_ref().join(format!("{}.{}", file_name, file_ext));
+    fn dump_info(&self, process: &mut IntoProcessInstanceArcBox<'_>) -> Result<()> {
+        let file_path = self.out_dir.join("info.json");
 
-        fs::write(&file_path, content)?;
-
-        Ok(())
-    }
-
-    fn dump_item<P: AsRef<Path>>(
-        &self,
-        item: &Item,
-        out_dir: P,
-        indent_size: usize,
-        file_exts: &[&str],
-        file_name: &str,
-    ) -> Result<()> {
-        for ext in file_exts {
-            let content = item.generate(self, indent_size, ext)?;
-
-            self.dump_file(&out_dir, file_name, ext, &content)?;
-        }
-
-        Ok(())
-    }
-
-    fn dump_info<P: AsRef<Path>>(
-        &self,
-        process: &mut IntoProcessInstanceArcBox<'_>,
-        out_dir: P,
-    ) -> Result<()> {
-        let content = &serde_json::to_string_pretty(&json!({
-            "timestamp": self.timestamp.to_rfc3339(),
-            "build_number": self.read_build_number(process).unwrap_or(0),
-        }))?;
-
-        self.dump_file(&out_dir, "info", "json", &content)
-    }
-
-    fn dump_schemas<P: AsRef<Path>>(
-        &self,
-        out_dir: P,
-        indent_size: usize,
-        file_exts: &[&str],
-    ) -> Result<()> {
-        for (module_name, (classes, enums)) in &self.schemas {
-            let map = SchemaMap::from([(module_name.clone(), (classes.clone(), enums.clone()))]);
-            let item = Item::Schemas(&map);
-
-            self.dump_item(&item, &out_dir, indent_size, file_exts, &module_name)?;
-        }
-
-        Ok(())
-    }
-
-    fn read_build_number(&self, process: &mut IntoProcessInstanceArcBox<'_>) -> Result<u32> {
-        self.offsets
+        let build_number = self
+            .result
+            .offsets
             .iter()
             .find_map(|(module_name, offsets)| {
                 let module = process.module_by_name(module_name).ok()?;
-                let offset = offsets.iter().find(|(name, _)| *name == "dwBuildNumber")?;
+                let offset = offsets.iter().find(|(name, _)| *name == "dwBuildNumber")?.1;
 
-                process.read(module.base + offset.1).ok()
+                process.read::<u32>(module.base + offset).ok()
             })
-            .ok_or(Error::Other("unable to read build number"))
+            .ok_or(Error::Other("unable to read build number"))?;
+
+        let content = serde_json::to_string_pretty(&json!({
+            "timestamp": self.timestamp.to_rfc3339(),
+            "build_number": build_number,
+        }))?;
+
+        fs::write(&file_path, &content)?;
+
+        Ok(())
+    }
+
+    fn dump_item(&self, file_name: &str, item: &Item) -> Result<()> {
+        for file_type in self.file_types {
+            let mut out = String::new();
+            let mut fmt = Formatter::new(&mut out, self.indent_size);
+
+            if file_type != "json" {
+                self.write_banner(&mut fmt)?;
+            }
+
+            item.write(&mut fmt, file_type)?;
+
+            let file_path = self.out_dir.join(format!("{}.{}", file_name, file_type));
+
+            fs::write(&file_path, out)?;
+        }
+
+        Ok(())
+    }
+
+    fn dump_schemas(&self) -> Result<()> {
+        for (module_name, (classes, enums)) in &self.result.schemas {
+            let map = SchemaMap::from([(module_name.clone(), (classes.clone(), enums.clone()))]);
+
+            self.dump_item(module_name, &Item::Schemas(&map))?;
+        }
+
+        Ok(())
     }
 
     fn write_banner(&self, fmt: &mut Formatter<'_>) -> Result<()> {
@@ -249,3 +193,7 @@ impl Results {
         Ok(())
     }
 }
+
+fn slugify(input: &str) -> String {
+    input.replace(|c: char| !c.is_alphanumeric(), "_")
+}