Add encrypted image support on sim

This adds new cargo features to allow running tests of encrypted
images with both RSA-OAEP and AES-128-KW.

When installing images on the simulated flash, both a plain and an
encrypted images are created. When encrypted image support is enabled,
verification of images in slot1 match against the encrypted image,
otherwise plain images are used.

PS: Also fixes ImageHeader to match bootutil definition.

Signed-off-by: Fabio Utzig <utzig@apache.org>
diff --git a/sim/src/lib.rs b/sim/src/lib.rs
index 1e8ba8a..e7ad671 100644
--- a/sim/src/lib.rs
+++ b/sim/src/lib.rs
@@ -1,5 +1,7 @@
 #[macro_use] extern crate log;
 extern crate ring;
+extern crate aes_ctr;
+extern crate base64;
 extern crate env_logger;
 extern crate docopt;
 extern crate libc;
@@ -18,6 +20,9 @@
 use std::mem;
 use std::process;
 use std::slice;
+use aes_ctr::Aes128Ctr;
+use aes_ctr::stream_cipher::generic_array::GenericArray;
+use aes_ctr::stream_cipher::{NewFixStreamCipher, StreamCipherCore};
 
 mod caps;
 mod tlv;
@@ -26,7 +31,7 @@
 use simflash::{Flash, SimFlash};
 use mcuboot_sys::{c, AreaDesc, FlashId};
 use caps::Caps;
-use tlv::TlvGen;
+use tlv::{TlvGen, TlvFlags, AES_SEC_KEY};
 
 const USAGE: &'static str = "
 Mcuboot simulator
@@ -183,12 +188,14 @@
         let slot0 = SlotInfo {
             base_off: slot0_base as usize,
             trailer_off: slot1_base - offset_from_end,
+            len: slot0_len as usize,
         };
 
         // And an upgrade image.
         let slot1 = SlotInfo {
             base_off: slot1_base as usize,
             trailer_off: scratch_base - offset_from_end,
+            len: slot1_len as usize,
         };
 
         Run {
@@ -216,15 +223,14 @@
     /// Construct an `Images` that doesn't expect an upgrade to happen.
     pub fn make_no_upgrade_image(&self) -> Images {
         let mut flash = self.flash.clone();
-        let primary = install_image(&mut flash, self.slots[0].base_off, 32784, false);
-        let upgrade = install_image(&mut flash, self.slots[1].base_off, 41928, false);
+        let primaries = install_image(&mut flash, &self.slots, 0, 32784, false);
+        let upgrades = install_image(&mut flash, &self.slots, 1, 41928, false);
         Images {
             flash: flash,
             areadesc: self.areadesc.clone(),
-            slot0: self.slots[0].clone(),
-            slot1: self.slots[1].clone(),
-            primary: primary,
-            upgrade: upgrade,
+            slots: [self.slots[0].clone(), self.slots[1].clone()],
+            primaries: primaries,
+            upgrades: upgrades,
             total_count: None,
             align: self.align,
             erased_val: self.erased_val,
@@ -234,7 +240,7 @@
     /// Construct an `Images` for normal testing.
     pub fn make_image(&self) -> Images {
         let mut images = self.make_no_upgrade_image();
-        mark_upgrade(&mut images.flash, &images.slot1);
+        mark_upgrade(&mut images.flash, &images.slots[1]);
 
         // upgrades without fails, counts number of flash operations
         let total_count = match images.run_basic_upgrade() {
@@ -250,15 +256,14 @@
 
     pub fn make_bad_slot1_image(&self) -> Images {
         let mut bad_flash = self.flash.clone();
-        let primary = install_image(&mut bad_flash, self.slots[0].base_off, 32784, false);
-        let upgrade = install_image(&mut bad_flash, self.slots[1].base_off, 41928, true);
+        let primaries = install_image(&mut bad_flash, &self.slots, 0, 32784, false);
+        let upgrades = install_image(&mut bad_flash, &self.slots, 1, 41928, true);
         Images {
             flash: bad_flash,
             areadesc: self.areadesc.clone(),
-            slot0: self.slots[0].clone(),
-            slot1: self.slots[1].clone(),
-            primary: primary,
-            upgrade: upgrade,
+            slots: [self.slots[0].clone(), self.slots[1].clone()],
+            primaries: primaries,
+            upgrades: upgrades,
             total_count: None,
             align: self.align,
             erased_val: self.erased_val,
@@ -380,7 +385,7 @@
         let (fl, total_count) = try_upgrade(&self.flash, &self, None);
         info!("Total flash operation count={}", total_count);
 
-        if !verify_image(&fl, self.slot0.base_off, &self.upgrade) {
+        if !verify_image(&fl, &self.slots, 0, &self.upgrades) {
             warn!("Image mismatch after first boot");
             Err(())
         } else {
@@ -402,7 +407,7 @@
             for count in 2 .. 5 {
                 info!("Try revert: {}", count);
                 let fl = try_revert(&self.flash, &self.areadesc, count, self.align);
-                if !verify_image(&fl, self.slot0.base_off, &self.primary) {
+                if !verify_image(&fl, &self.slots, 0, &self.primaries) {
                     error!("Revert failure on count {}", count);
                     fails += 1;
                 }
@@ -421,25 +426,25 @@
             info!("Try interruption at {}", i);
             let (fl, count) = try_upgrade(&self.flash, &self, Some(i));
             info!("Second boot, count={}", count);
-            if !verify_image(&fl, self.slot0.base_off, &self.upgrade) {
+            if !verify_image(&fl, &self.slots, 0, &self.upgrades) {
                 warn!("FAIL at step {} of {}", i, total_flash_ops);
                 fails += 1;
             }
 
-            if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+            if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                                BOOT_FLAG_SET, BOOT_FLAG_SET) {
                 warn!("Mismatched trailer for Slot 0");
                 fails += 1;
             }
 
-            if !verify_trailer(&fl, self.slot1.trailer_off, BOOT_MAGIC_UNSET,
+            if !verify_trailer(&fl, self.slots[1].trailer_off, BOOT_MAGIC_UNSET,
                                BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
                 warn!("Mismatched trailer for Slot 1");
                 fails += 1;
             }
 
             if Caps::SwapUpgrade.present() {
-                if !verify_image(&fl, self.slot1.base_off, &self.primary) {
+                if !verify_image(&fl, &self.slots, 1, &self.primaries) {
                     warn!("Slot 1 FAIL at step {} of {}", i, total_flash_ops);
                     fails += 1;
                 }
@@ -465,9 +470,9 @@
                                                   total_flash_ops, total_fails);
         info!("Random interruptions at reset points={:?}", total_counts);
 
-        let slot0_ok = verify_image(&fl, self.slot0.base_off, &self.upgrade);
+        let slot0_ok = verify_image(&fl, &self.slots, 0, &self.upgrades);
         let slot1_ok = if Caps::SwapUpgrade.present() {
-            verify_image(&fl, self.slot1.base_off, &self.primary)
+            verify_image(&fl, &self.slots, 1, &self.primaries)
         } else {
             true
         };
@@ -477,12 +482,12 @@
                    if slot1_ok { "ok" } else { "fail" });
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_SET, BOOT_FLAG_SET) {
             error!("Mismatched trailer for Slot 0");
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot1.trailer_off, BOOT_MAGIC_UNSET,
+        if !verify_trailer(&fl, self.slots[1].trailer_off, BOOT_MAGIC_UNSET,
                            BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
             error!("Mismatched trailer for Slot 1");
             fails += 1;
@@ -539,25 +544,25 @@
         //FIXME: copy_done is written by boot_go, is it ok if no copy
         //       was ever done?
 
-        if !verify_image(&fl, self.slot0.base_off, &self.upgrade) {
+        if !verify_image(&fl, &self.slots, 0, &self.upgrades) {
             warn!("Slot 0 image verification FAIL");
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_UNSET, BOOT_FLAG_SET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot1.trailer_off, BOOT_MAGIC_UNSET,
+        if !verify_trailer(&fl, self.slots[1].trailer_off, BOOT_MAGIC_UNSET,
                            BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
             warn!("Mismatched trailer for Slot 1");
             fails += 1;
         }
 
         // Marks image in slot0 as permanent, no revert should happen...
-        mark_permanent_upgrade(&mut fl, &self.slot0, self.align, self.erased_val);
+        mark_permanent_upgrade(&mut fl, &self.slots[0], self.align, self.erased_val);
 
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_SET, BOOT_FLAG_SET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
@@ -569,12 +574,12 @@
             fails += 1;
         }
 
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_SET, BOOT_FLAG_SET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
         }
-        if !verify_image(&fl, self.slot0.base_off, &self.upgrade) {
+        if !verify_image(&fl, &self.slots, 0, &self.upgrades) {
             warn!("Failed image verification");
             fails += 1;
         }
@@ -594,10 +599,10 @@
 
         info!("Try non-revert on imgtool generated image");
 
-        mark_upgrade(&mut fl, &self.slot0);
+        mark_upgrade(&mut fl, &self.slots[0]);
 
         // This simulates writing an image created by imgtool to Slot 0
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
@@ -611,16 +616,16 @@
         }
 
         // State should not have changed
-        if !verify_image(&fl, self.slot0.base_off, &self.primary) {
+        if !verify_image(&fl, &self.slots, 0, &self.primaries) {
             warn!("Failed image verification");
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot1.trailer_off, BOOT_MAGIC_UNSET,
+        if !verify_trailer(&fl, self.slots[1].trailer_off, BOOT_MAGIC_UNSET,
                            BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
             warn!("Mismatched trailer for Slot 1");
             fails += 1;
@@ -641,11 +646,11 @@
 
         info!("Try upgrade image with bad signature");
 
-        mark_upgrade(&mut fl, &self.slot0);
-        mark_permanent_upgrade(&mut fl, &self.slot0, self.align, self.erased_val);
-        mark_upgrade(&mut fl, &self.slot1);
+        mark_upgrade(&mut fl, &self.slots[0]);
+        mark_permanent_upgrade(&mut fl, &self.slots[0], self.align, self.erased_val);
+        mark_upgrade(&mut fl, &self.slots[1]);
 
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_SET, BOOT_FLAG_UNSET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
@@ -659,11 +664,11 @@
         }
 
         // State should not have changed
-        if !verify_image(&fl, self.slot0.base_off, &self.primary) {
+        if !verify_image(&fl, &self.slots, 0, &self.primaries) {
             warn!("Failed image verification");
             fails += 1;
         }
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_SET, BOOT_FLAG_UNSET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
@@ -683,10 +688,22 @@
 
     // FIXME: could get status sz from bootloader
     #[cfg(not(feature = "overwrite-only"))]
+    #[cfg(not(feature = "enc-rsa"))]
+    #[cfg(not(feature = "enc-kw"))]
     fn status_sz(&self) -> usize {
         self.trailer_sz() - (16 + 24)
     }
 
+    #[cfg(feature = "enc-rsa")]
+    fn status_sz(&self) -> usize {
+        self.trailer_sz() - (16 + 24 + 32)
+    }
+
+    #[cfg(feature = "enc-kw")]
+    fn status_sz(&self) -> usize {
+        self.trailer_sz() - (16 + 24 + 32)
+    }
+
     /// This test runs a simple upgrade with no fails in the images, but
     /// allowing for fails in the status area. This should run to the end
     /// and warn that write fails were detected...
@@ -700,9 +717,9 @@
 
         info!("Try swap with status fails");
 
-        mark_permanent_upgrade(&mut fl, &self.slot1, self.align, self.erased_val);
+        mark_permanent_upgrade(&mut fl, &self.slots[1], self.align, self.erased_val);
 
-        let status_off = self.slot1.base_off - self.trailer_sz();
+        let status_off = self.slots[1].base_off - self.trailer_sz();
 
         // Always fail writes to status area...
         let _ = fl.add_bad_region(status_off, self.status_sz(), 1.0);
@@ -720,13 +737,13 @@
             fails += 1;
         }
 
-        if !verify_trailer(&fl, self.slot0.trailer_off, BOOT_MAGIC_GOOD,
+        if !verify_trailer(&fl, self.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                            BOOT_FLAG_SET, BOOT_FLAG_SET) {
             warn!("Mismatched trailer for Slot 0");
             fails += 1;
         }
 
-        if !verify_image(&fl, self.slot0.base_off, &self.upgrade) {
+        if !verify_image(&fl, &self.slots, 0, &self.upgrades) {
             warn!("Failed image verification");
             fails += 1;
         }
@@ -758,9 +775,9 @@
 
         info!("Try interrupted swap with status fails");
 
-        mark_permanent_upgrade(&mut fl, &self.slot1, self.align, self.erased_val);
+        mark_permanent_upgrade(&mut fl, &self.slots[1], self.align, self.erased_val);
 
-        let status_off = self.slot1.base_off - self.trailer_sz();
+        let status_off = self.slots[1].base_off - self.trailer_sz();
 
         // Mark the status area as a bad area
         let _ = fl.add_bad_region(status_off, self.status_sz(), 0.5);
@@ -805,9 +822,9 @@
 
         info!("Try interrupted swap with status fails");
 
-        mark_permanent_upgrade(&mut fl, &self.slot1, self.align, self.erased_val);
+        mark_permanent_upgrade(&mut fl, &self.slots[1], self.align, self.erased_val);
 
-        let status_off = self.slot1.base_off - self.trailer_sz();
+        let status_off = self.slots[1].base_off - self.trailer_sz();
 
         // Mark the status area as a bad area
         let _ = fl.add_bad_region(status_off, self.status_sz(), 1.0);
@@ -835,7 +852,7 @@
     // Clone the flash to have a new copy.
     let mut fl = flash.clone();
 
-    mark_permanent_upgrade(&mut fl, &images.slot1, images.align, images.erased_val);
+    mark_permanent_upgrade(&mut fl, &images.slots[1], images.align, images.erased_val);
 
     let mut counter = stop.unwrap_or(0);
 
@@ -883,7 +900,7 @@
         fails += 1;
     }
 
-    if !verify_trailer(&fl, images.slot0.trailer_off, None, None, BOOT_FLAG_UNSET) {
+    if !verify_trailer(&fl, images.slots[0].trailer_off, None, None, BOOT_FLAG_UNSET) {
         warn!("copy_done should be unset");
         fails += 1;
     }
@@ -894,20 +911,20 @@
         fails += 1;
     }
 
-    if !verify_image(&fl, images.slot0.base_off, &images.upgrade) {
+    if !verify_image(&fl, &images.slots, 0, &images.upgrades) {
         warn!("Image in slot 0 before revert is invalid at stop={}", stop);
         fails += 1;
     }
-    if !verify_image(&fl, images.slot1.base_off, &images.primary) {
+    if !verify_image(&fl, &images.slots, 1, &images.primaries) {
         warn!("Image in slot 1 before revert is invalid at stop={}", stop);
         fails += 1;
     }
-    if !verify_trailer(&fl, images.slot0.trailer_off, BOOT_MAGIC_GOOD,
+    if !verify_trailer(&fl, images.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                        BOOT_FLAG_UNSET, BOOT_FLAG_SET) {
         warn!("Mismatched trailer for Slot 0 before revert");
         fails += 1;
     }
-    if !verify_trailer(&fl, images.slot1.trailer_off, BOOT_MAGIC_UNSET,
+    if !verify_trailer(&fl, images.slots[1].trailer_off, BOOT_MAGIC_UNSET,
                        BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
         warn!("Mismatched trailer for Slot 1 before revert");
         fails += 1;
@@ -920,20 +937,20 @@
         fails += 1;
     }
 
-    if !verify_image(&fl, images.slot0.base_off, &images.primary) {
+    if !verify_image(&fl, &images.slots, 0, &images.primaries) {
         warn!("Image in slot 0 after revert is invalid at stop={}", stop);
         fails += 1;
     }
-    if !verify_image(&fl, images.slot1.base_off, &images.upgrade) {
+    if !verify_image(&fl, &images.slots, 1, &images.upgrades) {
         warn!("Image in slot 1 after revert is invalid at stop={}", stop);
         fails += 1;
     }
-    if !verify_trailer(&fl, images.slot0.trailer_off, BOOT_MAGIC_GOOD,
+    if !verify_trailer(&fl, images.slots[0].trailer_off, BOOT_MAGIC_GOOD,
                        BOOT_FLAG_SET, BOOT_FLAG_SET) {
         warn!("Mismatched trailer for Slot 1 after revert");
         fails += 1;
     }
-    if !verify_trailer(&fl, images.slot1.trailer_off, BOOT_MAGIC_UNSET,
+    if !verify_trailer(&fl, images.slots[1].trailer_off, BOOT_MAGIC_UNSET,
                        BOOT_FLAG_UNSET, BOOT_FLAG_UNSET) {
         warn!("Mismatched trailer for Slot 1 after revert");
         fails += 1;
@@ -946,7 +963,7 @@
                     total_ops: i32,  count: usize) -> (SimFlash, Vec<i32>) {
     let mut fl = flash.clone();
 
-    mark_permanent_upgrade(&mut fl, &images.slot1, images.align, images.erased_val);
+    mark_permanent_upgrade(&mut fl, &images.slots[1], images.align, images.erased_val);
 
     let mut rng = rand::thread_rng();
     let mut resets = vec![0i32; count];
@@ -985,20 +1002,19 @@
 
 /// Install a "program" into the given image.  This fakes the image header, or at least all of the
 /// fields used by the given code.  Returns a copy of the image that was written.
-fn install_image(flash: &mut Flash, offset: usize, len: usize,
-                 bad_sig: bool) -> Vec<u8> {
-    let offset0 = offset;
+fn install_image(flash: &mut Flash, slots: &[SlotInfo], slot: usize, len: usize,
+                 bad_sig: bool) -> [Option<Vec<u8>>; 2] {
+    let offset = slots[slot].base_off;
+    let slot_len = slots[slot].len;
 
     let mut tlv = make_tlv();
 
     // Generate a boot header.  Note that the size doesn't include the header.
     let header = ImageHeader {
         magic: 0x96f3b83d,
-        tlv_size: tlv.get_size(),
-        _pad1: 0,
+        load_addr: 0,
         hdr_size: 32,
-        key_id: 0,
-        _pad2: 0,
+        _pad1: 0,
         img_size: len as u32,
         flags: tlv.get_flags(),
         ver: ImageVersion {
@@ -1007,7 +1023,7 @@
             revision: 1,
             build_num: offset as u32,
         },
-        _pad3: 0,
+        _pad2: 0,
     };
 
     let b_header = header.as_raw();
@@ -1017,35 +1033,106 @@
                                                   mem::size_of::<ImageHeader>()) };
                                                   */
     assert_eq!(b_header.len(), 32);
-    flash.write(offset, &b_header).unwrap();
-    let offset = offset + b_header.len();
 
     // The core of the image itself is just pseudorandom data.
-    let mut buf = vec![0; len];
-    splat(&mut buf, offset);
-    tlv.add_bytes(&buf);
+    let mut b_img = vec![0; len];
+    splat(&mut b_img, offset);
 
-    // Get and append the TLV itself.
-    if bad_sig {
-        let good_sig = &mut tlv.make_tlv();
-        buf.append(&mut vec![0; good_sig.len()]);
-    } else {
-        buf.append(&mut tlv.make_tlv());
+    // TLV signatures work over plain image
+    tlv.add_bytes(&b_img);
+
+    // Generate encrypted images
+    let flag = TlvFlags::ENCRYPTED as u32;
+    let is_encrypted = (tlv.get_flags() & flag) == flag;
+    let mut b_encimg = vec![];
+    if is_encrypted {
+        let key = GenericArray::from_slice(AES_SEC_KEY);
+        let nonce = GenericArray::from_slice(&[0; 16]);
+        let mut cipher = Aes128Ctr::new(&key, &nonce);
+        b_encimg = b_img.clone();
+        cipher.apply_keystream(&mut b_encimg);
     }
 
+    // Build the TLV itself.
+    let mut b_tlv = if bad_sig {
+        let good_sig = &mut tlv.make_tlv();
+        vec![0; good_sig.len()]
+    } else {
+        tlv.make_tlv()
+    };
+
     // Pad the block to a flash alignment (8 bytes).
-    while buf.len() % 8 != 0 {
-        buf.push(0xFF);
+    while b_tlv.len() % 8 != 0 {
+        //FIXME: should be erase_val?
+        b_tlv.push(0xFF);
     }
 
-    flash.write(offset, &buf).unwrap();
-    let offset = offset + buf.len();
+    let mut buf = vec![];
+    buf.append(&mut b_header.to_vec());
+    buf.append(&mut b_img);
+    buf.append(&mut b_tlv.clone());
 
-    // Copy out the image so that we can verify that the image was installed correctly later.
-    let mut copy = vec![0u8; offset - offset0];
-    flash.read(offset0, &mut copy).unwrap();
+    let mut encbuf = vec![];
+    if is_encrypted {
+        encbuf.append(&mut b_header.to_vec());
+        encbuf.append(&mut b_encimg);
+        encbuf.append(&mut b_tlv);
+    }
 
-    copy
+    let result: [Option<Vec<u8>>; 2];
+
+    // Since images are always non-encrypted in slot0, we first write an
+    // encrypted image, re-read to use for verification, erase + flash
+    // un-encrypted. In slot1 the image is written un-encrypted, and if
+    // encryption is requested, it follows an erase + flash encrypted.
+
+    if slot == 0 {
+        let enc_copy: Option<Vec<u8>>;
+
+        if is_encrypted {
+            flash.write(offset, &encbuf).unwrap();
+
+            let mut enc = vec![0u8; encbuf.len()];
+            flash.read(offset, &mut enc).unwrap();
+
+            enc_copy = Some(enc);
+
+            flash.erase(offset, slot_len).unwrap();
+        } else {
+            enc_copy = None;
+        }
+
+        flash.write(offset, &buf).unwrap();
+
+        let mut copy = vec![0u8; buf.len()];
+        flash.read(offset, &mut copy).unwrap();
+
+        result = [Some(copy), enc_copy];
+    } else {
+        flash.write(offset, &buf).unwrap();
+
+        let mut copy = vec![0u8; buf.len()];
+        flash.read(offset, &mut copy).unwrap();
+
+        let enc_copy: Option<Vec<u8>>;
+
+        if is_encrypted {
+            flash.erase(offset, slot_len).unwrap();
+
+            flash.write(offset, &encbuf).unwrap();
+
+            let mut enc = vec![0u8; encbuf.len()];
+            flash.read(offset, &mut enc).unwrap();
+
+            enc_copy = Some(enc);
+        } else {
+            enc_copy = None;
+        }
+
+        result = [Some(copy), enc_copy];
+    }
+
+    result
 }
 
 // The TLV in use depends on what kind of signature we are verifying.
@@ -1059,21 +1146,64 @@
     TlvGen::new_ecdsa()
 }
 
+#[cfg(feature = "enc-rsa")]
+fn make_tlv() -> TlvGen {
+    TlvGen::new_enc_rsa()
+}
+
+#[cfg(feature = "enc-kw")]
+fn make_tlv() -> TlvGen {
+    TlvGen::new_enc_kw()
+}
+
 #[cfg(not(feature = "sig-rsa"))]
 #[cfg(not(feature = "sig-ecdsa"))]
+#[cfg(not(feature = "enc-rsa"))]
+#[cfg(not(feature = "enc-kw"))]
 fn make_tlv() -> TlvGen {
     TlvGen::new_hash_only()
 }
 
+#[cfg(feature = "enc-rsa")]
+fn find_image(images: &[Option<Vec<u8>>; 2], slot: usize) -> &Vec<u8> {
+    match &images[slot] {
+        Some(image) => return image,
+        None => panic!("Invalid image"),
+    }
+}
+
+#[cfg(feature = "enc-kw")]
+fn find_image(images: &[Option<Vec<u8>>; 2], slot: usize) -> &Vec<u8> {
+    match &images[slot] {
+        Some(image) => return image,
+        None => panic!("Invalid image"),
+    }
+}
+
+#[cfg(not(feature = "enc-rsa"))]
+#[cfg(not(feature = "enc-kw"))]
+fn find_image(images: &[Option<Vec<u8>>; 2], _slot: usize) -> &Vec<u8> {
+    match &images[0] {
+        Some(image) => return image,
+        None => panic!("Invalid image"),
+    }
+}
+
 /// Verify that given image is present in the flash at the given offset.
-fn verify_image(flash: &Flash, offset: usize, buf: &[u8]) -> bool {
+fn verify_image(flash: &Flash, slots: &[SlotInfo], slot: usize,
+                images: &[Option<Vec<u8>>; 2]) -> bool {
+    let image = find_image(images, slot);
+    let buf = image.as_slice();
+
     let mut copy = vec![0u8; buf.len()];
+    let offset = slots[slot].base_off;
     flash.read(offset, &mut copy).unwrap();
 
     if buf != &copy[..] {
         for i in 0 .. buf.len() {
             if buf[i] != copy[i] {
-                info!("First failure at {:#x}", offset + i);
+                info!("First failure for slot{} at {:#x} {:#x}!={:#x}",
+                      slot, offset + i, buf[i], copy[i]);
                 break;
             }
         }
@@ -1153,15 +1283,13 @@
 #[repr(C)]
 pub struct ImageHeader {
     magic: u32,
-    tlv_size: u16,
-    key_id: u8,
-    _pad1: u8,
+    load_addr: u32,
     hdr_size: u16,
-    _pad2: u16,
+    _pad1: u16,
     img_size: u32,
     flags: u32,
     ver: ImageVersion,
-    _pad3: u32,
+    _pad2: u32,
 }
 
 impl AsRaw for ImageHeader {}
@@ -1178,15 +1306,15 @@
 struct SlotInfo {
     base_off: usize,
     trailer_off: usize,
+    len: usize,
 }
 
 pub struct Images {
     flash: SimFlash,
     areadesc: AreaDesc,
-    slot0: SlotInfo,
-    slot1: SlotInfo,
-    primary: Vec<u8>,
-    upgrade: Vec<u8>,
+    slots: [SlotInfo; 2],
+    primaries: [Option<Vec<u8>>; 2],
+    upgrades: [Option<Vec<u8>>; 2],
     total_count: Option<i32>,
     align: u8,
     erased_val: u8,
diff --git a/sim/src/tlv.rs b/sim/src/tlv.rs
index 79e89fd..7878dd2 100644
--- a/sim/src/tlv.rs
+++ b/sim/src/tlv.rs
@@ -10,6 +10,7 @@
 
 use std::sync::Arc;
 use pem;
+use base64;
 use ring::{digest, rand, signature};
 use untrusted;
 use mcuboot_sys::c;
@@ -23,6 +24,16 @@
     RSA2048 = 0x20,
     ECDSA224 = 0x21,
     ECDSA256 = 0x22,
+    ENCRSA2048 = 0x30,
+    ENCKW128 = 0x31,
+}
+
+#[allow(dead_code, non_camel_case_types)]
+pub enum TlvFlags {
+    PIC = 0x01,
+    NON_BOOTABLE = 0x02,
+    ENCRYPTED = 0x04,
+    RAM_LOAD = 0x20,
 }
 
 pub struct TlvGen {
@@ -32,6 +43,8 @@
     payload: Vec<u8>,
 }
 
+pub const AES_SEC_KEY: &[u8; 16] = b"0123456789ABCDEF";
+
 impl TlvGen {
     /// Construct a new tlv generator that will only contain a hash of the data.
     #[allow(dead_code)]
@@ -64,6 +77,26 @@
         }
     }
 
+    #[allow(dead_code)]
+    pub fn new_enc_rsa() -> TlvGen {
+        TlvGen {
+            flags: TlvFlags::ENCRYPTED as u32,
+            kinds: vec![TlvKinds::SHA256, TlvKinds::ENCRSA2048],
+            size: 4 + 32 + 4 + 256,
+            payload: vec![],
+        }
+    }
+
+    #[allow(dead_code)]
+    pub fn new_enc_kw() -> TlvGen {
+        TlvGen {
+            flags: TlvFlags::ENCRYPTED as u32,
+            kinds: vec![TlvKinds::SHA256, TlvKinds::ENCKW128],
+            size: 4 + 32 + 4 + 24,
+            payload: vec![],
+        }
+    }
+
     /// Retrieve the header flags for this configuration.  This can be called at any time.
     pub fn get_flags(&self) -> u32 {
         self.flags
@@ -212,6 +245,41 @@
             result.extend(der);
         }
 
+        if self.kinds.contains(&TlvKinds::ENCRSA2048) {
+            let key_bytes = pem::parse(include_bytes!("../../enc-rsa2048-pub.pem")
+                                       .as_ref()).unwrap();
+            assert_eq!(key_bytes.tag, "PUBLIC KEY");
+
+            let encbuf = match c::rsa_oaep_encrypt(&key_bytes.contents, AES_SEC_KEY) {
+                Ok(v) => v,
+                Err(_) => panic!("Failed to encrypt secret key"),
+            };
+
+            assert!(encbuf.len() == 256);
+            result.push(TlvKinds::ENCRSA2048 as u8);
+            result.push(0);
+            result.push(0);
+            result.push(1);
+            result.extend_from_slice(&encbuf);
+        }
+
+        if self.kinds.contains(&TlvKinds::ENCKW128) {
+            let key_bytes = base64::decode(
+                include_str!("../../enc-aes128kw.b64").trim()).unwrap();
+
+            let encbuf = match c::kw_encrypt(&key_bytes, AES_SEC_KEY) {
+                Ok(v) => v,
+                Err(_) => panic!("Failed to encrypt secret key"),
+            };
+
+            assert!(encbuf.len() == 24);
+            result.push(TlvKinds::ENCKW128 as u8);
+            result.push(0);
+            result.push(24);
+            result.push(0);
+            result.extend_from_slice(&encbuf);
+        }
+
         result
     }
 }