Add keygen command

MegatonHammer · Jan 25, 2019 · 4c57e2c · 4c57e2c
1 parent 570a418
commit 4c57e2c
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 3 deletions.
diff --git a/src/bin/linkle_clap.rs b/src/bin/linkle_clap.rs
@@ -74,6 +74,17 @@ enum Opt {
         #[structopt(parse(from_os_str))]
         output_file: PathBuf,
     },
+    /// Print all the keys generated from our keyfile.
+    #[structopt(name = "keygen")]
+    Keygen {
+        /// Use development keys instead of retail
+        #[structopt(short = "d", long = "dev")]
+        dev: bool,
+
+        /// Key file to use
+        #[structopt(parse(from_os_str), short = "k", long = "keyset")]
+        keyfile: Option<PathBuf>,
+    }
 }
 
 fn create_nxo(format: &str, input_file: &str, output_file: &str, icon_file: Option<&str>, romfs_dir: Option<&str>, nacp_file: Option<&str>) -> Result<(), linkle::error::Error> {
@@ -153,6 +164,13 @@ fn create_romfs(input_directory: &Path, output_file: &Path) -> Result<(), linkle
     Ok(())
 }
 
+fn print_keys(is_dev: bool, key_path: Option<&Path>) -> Result<(), linkle::error::Error> {
+    let keys = linkle::pki::Keys::new_retail(key_path).unwrap();
+
+    keys.write(&mut std::io::stdout()).unwrap();
+    Ok(())
+}
+
 fn to_opt_ref<U: ?Sized, T: AsRef<U>>(s: &Option<T>) -> Option<&U> {
     s.as_ref().map(AsRef::as_ref)
 }
@@ -165,6 +183,7 @@ fn process_args(app: &Opt) {
         Opt::Pfs0Extract { ref input_file, ref output_directory } => extract_pfs0(input_file, output_directory),
         Opt::Nacp { ref input_file, ref output_file } => create_nacp(input_file, output_file),
         Opt::Romfs { ref input_directory, ref output_file } => create_romfs(input_directory, output_file),
+        Opt::Keygen { dev, ref keyfile } => print_keys(*dev, to_opt_ref(keyfile)),
     };
 
     if let Err(e) = res {

diff --git a/src/pki.rs b/src/pki.rs
@@ -2,7 +2,7 @@ use std::fmt;
 use ini::{self, ini::Properties};
 use failure::Backtrace;
 use std::fs::File;
-use std::io::{self, ErrorKind};
+use std::io::{self, Write, ErrorKind};
 use std::path::Path;
 use crate::error::Error;
 use aes::Aes128;
@@ -29,6 +29,14 @@ macro_rules! impl_debug {
                 Ok(())
             }
         }
+        impl fmt::Display for $for {
+            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+                for byte in &self.0[..] {
+                    write!(f, "{:02X}", byte)?;
+                }
+                Ok(())
+            }
+        }
     }
 }
 
@@ -171,6 +179,71 @@ pub struct Keys {
     package2_fixed_key_modulus: Option<Modulus>,
 }
 
+macro_rules! make_key_macros_write {
+    ($self:ident, $w:ident) => {
+        macro_rules! single_key {
+            ($keyname:tt) => {
+                if let Some(key) = &$self.$keyname {
+                    writeln!($w, "{} = {}", stringify!($keyname), key)?;
+                }
+            }
+        }
+
+        macro_rules! single_key_xts {
+            ($keyname:tt) => {
+                if let Some(key) = &$self.$keyname {
+                    writeln!($w, "{} = {}", stringify!($keyname), key)?;
+                }
+            }
+        }
+
+        macro_rules! multi_key {
+            ($keyname:tt) => {
+                for (idx, v) in $self.$keyname.iter().enumerate() {
+                    if let Some(key) = v {
+                        // remove trailing s
+                        let mut name = String::from(stringify!($keyname));
+                        if name.bytes().last() == Some(b's') {
+                            name.pop();
+                        }
+                        writeln!($w, "{}_{:02x} = {}", name, idx, key)?;
+                    }
+                }
+            }
+        }
+
+        macro_rules! multi_keyblob {
+            ($keyname:tt) => {
+                for (idx, v) in $self.$keyname.iter().enumerate() {
+                    if let Some(key) = v {
+                        // remove trailing s
+                        let mut name = String::from(stringify!($keyname));
+                        if name.bytes().last() == Some(b's') {
+                            name.pop();
+                        }
+                        writeln!($w, "{}_{:02x} = {}", name, idx, key)?;
+                    }
+                }
+            }
+        }
+
+        macro_rules! multi_encrypted_keyblob {
+            ($keyname:tt) => {
+                for (idx, v) in $self.$keyname.iter().enumerate() {
+                    if let Some(key) = v {
+                        // remove trailing s
+                        let mut name = String::from(stringify!($keyname));
+                        if name.bytes().last() == Some(b's') {
+                            name.pop();
+                        }
+                        writeln!($w, "{}_{:02x} = {}", name, idx, key)?;
+                    }
+                }
+            }
+        }
+    }
+}
+
 macro_rules! make_key_macros {
     ($self:ident, $section:ident) => {
         macro_rules! single_key {
@@ -193,7 +266,9 @@ macro_rules! make_key_macros {
                     let mut key = [0; 0x10];
                     // remove trailing s
                     let mut name = String::from(stringify!($keyname));
-                    name.pop();
+                    if name.bytes().last() == Some(b's') {
+                        name.pop();
+                    }
                     v.or_in(key_to_aes_array($section, &name, idx, &mut key)?.map(|()| Aes128Key(key)));
                 }
             }
@@ -219,7 +294,9 @@ macro_rules! make_key_macros {
                     let mut key = [0; 0xB0];
                     // remove trailing s
                     let mut name = String::from(stringify!($keyname));
-                    name.pop();
+                    if name.bytes().last() == Some(b's') {
+                        name.pop();
+                    }
                     v.or_in(key_to_aes_array($section, &name, idx, &mut key)?.map(|()| EncryptedKeyblob(key)));
                 }
             }
@@ -380,6 +457,47 @@ impl Keys {
         Ok(())
     }
 
+    pub fn write<W: Write>(&self, w: &mut W) -> io::Result<()> {
+        make_key_macros_write!(self, w);
+        single_key!(secure_boot_key);
+        single_key!(tsec_key);
+        multi_key!(keyblob_keys);
+        multi_key!(keyblob_mac_keys);
+        multi_key!(keyblob_key_sources);
+        multi_encrypted_keyblob!(encrypted_keyblobs);
+        multi_keyblob!(keyblobs);
+        single_key!(keyblob_mac_key_source);
+        single_key!(tsec_root_key);
+        multi_key!(master_kek_sources);
+        multi_key!(master_keks);
+        single_key!(master_key_source);
+        multi_key!(master_keys);
+        multi_key!(package1_keys);
+        multi_key!(package2_keys);
+        single_key!(package2_key_source);
+        single_key!(aes_kek_generation_source);
+        single_key!(aes_key_generation_source);
+        single_key!(key_area_key_application_source);
+        single_key!(key_area_key_ocean_source);
+        single_key!(key_area_key_system_source);
+        single_key!(titlekek_source);
+        single_key!(header_kek_source);
+        single_key!(sd_card_kek_source);
+        single_key_xts!(sd_card_save_key_source);
+        single_key_xts!(sd_card_nca_key_source);
+        single_key!(save_mac_kek_source);
+        single_key!(save_mac_key_source);
+        single_key_xts!(header_key_source);
+        single_key_xts!(header_key);
+        multi_key!(titlekeks);
+        multi_key!(key_area_key_application);
+        multi_key!(key_area_key_ocean);
+        multi_key!(key_area_key_system);
+        single_key_xts!(sd_card_save_key);
+        single_key_xts!(sd_card_nca_key);
+        Ok(())
+    }
+
     pub fn derive_keys(&mut self) -> Result<(), Error> {
         for i in 0..6 {
             /* Derive the keyblob_keys */