From 406ed1325f894b6606e8626b5ca03dab0ef042de Mon Sep 17 00:00:00 2001
From: Roey Darwish Dror <roey.ghost@gmail.com>
Date: Wed, 7 Nov 2018 11:09:18 +0200
Subject: [PATCH] Add chroot subcommand (fix #13)

---
 README.md   |  7 +++++++
 src/main.rs | 55 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/README.md b/README.md
index 182c432..1f32362 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,7 @@ Dependencies will be handled for you if you install alma from AUR.
 
 ## Usage
 
+### Creation
 ``` shell
 sudo alma create /dev/disk/by-id/usb-Generic_USB_Flash_Disk-0:0
 ```
@@ -32,6 +33,12 @@ this tool will refuse to work with drive paths which don't start with `/dev/disk
 After the installation is done you can either boot from it immediately or use `arch-chroot` to
 perform further customizations before your first boot.
 
+### Chrooting to exiting Live USB
+
+``` shell
+sudo alma chroot /dev/disk/by-id/usb-Generic_USB_Flash_Disk-0:0
+```
+
 ### Flags
 * `-p / --extra-packages` - Specify extra packages to install. For example: `-p htop tmux`
 * `-i / --interactive` - Drop you into interactive chroot to perform further customization
diff --git a/src/main.rs b/src/main.rs
index dc80611..e778036 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -43,6 +43,9 @@ SystemMaxUse=16M
 enum App {
     #[structopt(name = "create", about = "Create a new Arch Linux USB")]
     Create(CreateCommand),
+
+    #[structopt(name = "chroot", about = "Chroot into exiting Live USB")]
+    Chroot(ChrootCommand),
 }
 
 #[derive(StructOpt)]
@@ -60,6 +63,13 @@ struct CreateCommand {
     interactive: bool,
 }
 
+#[derive(StructOpt)]
+struct ChrootCommand {
+    /// Path starting with /dev/disk/by-id for the USB drive
+    #[structopt(parse(from_os_str),)]
+    disk: PathBuf,
+}
+
 fn create(command: CreateCommand) -> Result<(), Error> {
     let sgdisk = Tool::find("sgdisk")?;
     let pacstrap = Tool::find("pacstrap")?;
@@ -196,6 +206,50 @@ fn create(command: CreateCommand) -> Result<(), Error> {
     Ok(())
 }
 
+fn chroot(command: ChrootCommand) -> Result<(), Error> {
+    let arch_chroot = Tool::find("arch-chroot")?;
+
+    if !(command.disk.starts_with("/dev/disk/by-id")
+        && (command
+            .disk
+            .file_name()
+            .and_then(|s| s.to_str())
+            .filter(|ref f| f.starts_with("usb-"))
+            .is_some()))
+    {
+        Err(ErrorKind::NotUSB)?;
+    }
+
+    let mount_point = tempdir().context(ErrorKind::TmpDirError)?;
+    let boot_point = mount_point.path().join("boot");
+    let mut mount_stack = MountStack::new();
+
+    info!("Mounting filesystems to {}", mount_point.path().display());
+    mount_stack
+        .mount(
+            &PathBuf::from(format!("{}-part3", command.disk.display())),
+            &mount_point.path(),
+            Filesystem::Btrfs,
+            Some("compress=zstd"),
+        ).context(ErrorKind::Mounting)?;
+
+    mount_stack
+        .mount(
+            &PathBuf::from(format!("{}-part2", command.disk.display())),
+            &boot_point,
+            Filesystem::Vfat,
+            None,
+        ).context(ErrorKind::Mounting)?;
+
+    arch_chroot
+        .execute()
+        .arg(mount_point.path())
+        .run(ErrorKind::Interactive)?;
+
+    info!("Unmounting filesystems");
+    Ok(())
+}
+
 extern "C" fn handle_sigint(_: i32) {
     warn!("Interrupted");
 }
@@ -220,6 +274,7 @@ fn main() {
 
     let result = match app {
         App::Create(command) => create(command),
+        App::Chroot(command) => chroot(command),
     };
 
     match result {