From 5f23997252e9bca08e8045b76498dde536124f07 Mon Sep 17 00:00:00 2001
From: m0veax <patrick.kilter@gmail.com>
Date: Sat, 27 Jan 2024 23:19:02 +0100
Subject: [PATCH] add extract-dtb command

---
 .gitignore          |  1 +
 Cargo.toml          |  5 ++++
 extract-dtb/main.rs | 67 +++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs          | 16 +++--------
 4 files changed, 77 insertions(+), 12 deletions(-)
 create mode 100644 extract-dtb/main.rs

diff --git a/.gitignore b/.gitignore
index 2f7896d..2c96eb1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 target/
+Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
index 6c65c4f..e3bb549 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,6 +12,11 @@ documentation = "https://docs.rs/psi_device_tree"
 hashbrown = "0.13"
 serde_json = "1.0"
 serde = { version = "1.0", features = ["derive"] }
+clap = { version = "4.4.18", features = ["derive"] }
 
 [features]
 string-dedup = [] # Requires std
+
+[[bin]]
+name = "extract-dtb"
+path = "extract-dtb/main.rs"
diff --git a/extract-dtb/main.rs b/extract-dtb/main.rs
new file mode 100644
index 0000000..2c48216
--- /dev/null
+++ b/extract-dtb/main.rs
@@ -0,0 +1,67 @@
+use clap::Parser;
+use psi_device_tree::DeviceTree as DT;
+use std::fs::{self, File};
+use std::io::{Read, Seek, SeekFrom, Write};
+use std::os::unix::fs::MetadataExt;
+
+// doodfeed is not a burger
+const DTB_MAGIC: u32 = 0xd00d_feed;
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Filename
+    #[arg(short, long)]
+    filename: String,
+    #[arg(short, long)]
+    dest: String,
+}
+
+fn main() {
+    let args = Args::parse();
+    let filename = args.filename;
+    let dest = args.dest;
+
+    let mut f = File::open(filename).unwrap();
+    let step = 8;
+    let size = f.metadata().unwrap().size();
+
+    for o in (0..size).step_by(step as usize) {
+        // read first bytes
+        f.seek(SeekFrom::Start(o)).unwrap();
+        let buf = &mut [0u8; 4];
+        f.read(buf).unwrap();
+
+        // is le magic?
+        if u32::from_be_bytes(*buf) == DTB_MAGIC {
+            // next 4 bytes plz to get size of device tree
+            f.read(buf).unwrap();
+            // size is not little endian
+            let size = u32::from_be_bytes(*buf) as usize;
+
+            f.seek(SeekFrom::Start(o)).unwrap();
+            // create vec of size filled with zerozero
+            let mut buf = vec![0; size];
+            f.read_exact(&mut buf).unwrap();
+
+            let dt = DT::load(&buf).unwrap();
+
+            // does it exist? if not, rule 34
+            fs::create_dir_all(&dest).unwrap();
+
+            let mut filename = String::from(format!("{o:08x}"));
+            filename.push_str(".dtb");
+
+            let path = std::path::Path::new(&dest).join(filename);
+
+            println!("Write dtb to {}", path.display());
+
+            let mut output = std::fs::OpenOptions::new()
+                .write(true)
+                .create(true)
+                .open(path)
+                .unwrap();
+            output.write_all(&buf).unwrap();
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index cd01ea8..26a8788 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -48,10 +48,10 @@ use serde::{Deserialize, Serialize};
 pub use error::*;
 use util::{align, SliceRead, VecWrite};
 
-#[cfg(not(feature = "string-dedup"))]
-mod string_table;
 #[cfg(feature = "string-dedup")]
 mod advanced_string_table;
+#[cfg(not(feature = "string-dedup"))]
+mod string_table;
 
 #[cfg(not(feature = "string-dedup"))]
 use string_table::StringTable;
@@ -247,11 +247,7 @@ impl DeviceTree {
 }
 
 impl Node {
-    fn load(
-        buffer: &[u8],
-        start: usize,
-        off_dt_strings: usize,
-    ) -> Result<(usize, Node)> {
+    fn load(buffer: &[u8], start: usize, off_dt_strings: usize) -> Result<(usize, Node)> {
         // check for DT_BEGIN_NODE
         if buffer.read_be_u32(start)? != OF_DT_BEGIN_NODE {
             return Err(Error::ParseError(start));
@@ -370,11 +366,7 @@ impl Node {
         Ok(raw.as_slice().read_be_u32(0)?)
     }
 
-    pub fn store(
-        &self,
-        structure: &mut Vec<u8>,
-        strings: &mut StringTable,
-    ) -> Result<()> {
+    pub fn store(&self, structure: &mut Vec<u8>, strings: &mut StringTable) -> Result<()> {
         structure.pad(4)?;
         let len = structure.len();
         structure.write_be_u32(len, OF_DT_BEGIN_NODE)?;