refactor + firewall

src/config.rs

@@ -88,14 +88,16 @@ pub struct GeneralConfig {
     // Enable Bluetooth
     pub bluetooth: Option<bool>,
     /// Install Video Driver
-    pub gpu_driver: Option<GPUVendor>
+    pub gpu_driver: Option<GPUVendor>,
+    // Eanble firewall
+    pub firewall: Option<bool>,
 }
 
 #[derive(Debug, Clone, Deserialize)]
 pub enum GPUVendor {
     AMD,
     NVIDIA,
-    INTEL
+    INTEL,
 }
 
 #[derive(Debug, Clone, Deserialize)]

src/install/drives.rs

@@ -1,10 +1,13 @@
-use crate::{config::DriveConfig, linux::run_command};
+use yansi::{Color, Paint};
+
+use crate::{config::DriveConfig, linux::run_command, print_status};
 
 // TODO : Add support for using entire block device
 
 /// Format the drives with the given config
 pub fn format_drives(conf: &DriveConfig) {
-    // TODO : Safety checks !!!
+    disk_safe_check(&conf.root);
+    disk_safe_check(&conf.root);
 
     // EFI (BOOT)
     run_command(&["mkfs.vfat", "-F", "32", conf.boot.as_str()], None, false);
@@ -52,3 +55,59 @@ pub fn mount_drives(conf: &DriveConfig) {
         false,
     );
 }
+
+/// Format a disk with EFI and BOOT partitions.
+pub fn partition_disk(dev: &str) {
+    disk_safe_check(dev);
+    print_status(&format!("Formatting disk {dev}"));
+    let cmd = vec!["fdisk", dev];
+    let input = "g\nn\n1\n\n+1G\nt\n1\nn\n2\n\nw\n";
+    run_command(&cmd, Some(input), false);
+}
+
+/// Check if `dev` contains a filesystem and error if it does
+pub fn disk_safe_check(dev: &str) {
+    if let Some(fs) = has_filesystem(dev) {
+        println!(
+            "{} Device {} already contains a filesystem {}",
+            "Error:".paint(Color::Red),
+            dev.paint(Color::Red),
+            fs.paint(Color::Blue)
+        );
+        std::process::exit(1);
+    }
+}
+
+pub fn has_filesystem(dev: &str) -> Option<String> {
+    let blockdevs: serde_json::Value =
+        serde_json::from_str(&run_command(&["lsblk", "--fs", "--json"], None, false).0).unwrap();
+
+    // TODO : Follow if symlink from /dev/disks
+    let dev = dev.trim_start_matches("/dev/");
+
+    let dev_entry = blockdevs
+        .as_object()
+        .unwrap()
+        .get("blockdevices")
+        .unwrap()
+        .as_array()
+        .unwrap()
+        .iter()
+        .find(|x| {
+            x.as_object()
+                .unwrap()
+                .get("name")
+                .unwrap()
+                .as_str()
+                .unwrap()
+                == dev
+        });
+
+    if let Some(dev_entry) = dev_entry {
+        if let Some(fs) = dev_entry.as_object().unwrap().get("fstype") {
+            return Some(fs.to_string());
+        }
+    }
+
+    None
+}

src/install/firewall.rs

@@ -0,0 +1,21 @@
+use crate::{
+    linux::{arch_chroot, systemd_service_enable},
+    pkg::install_pkgs,
+    print_status,
+};
+
+pub fn setup_firewall(ssh: bool) {
+    print_status("Enabling firewall");
+    install_pkgs(&["ufw"]);
+    systemd_service_enable("ufw.service");
+
+    arch_chroot(&["ufw", "default", "deny"], None, false);
+
+    if ssh {
+        arch_chroot(&["ufw", "allow", "22/tcp"], None, false);
+
+        arch_chroot(&["ufw", "limit", "ssh"], None, false);
+    }
+
+    arch_chroot(&["ufw", "enable"], None, false);
+}

src/install/first_boot.rs

@@ -45,7 +45,14 @@ pub fn first_boot_values(conf: &GeneralConfig) {
     print_status("Writing /etc/hostname");
     std::fs::write("/mnt/etc/hostname", format!("{}\n", conf.hostname)).unwrap();
 
-    install_file("/mnt/etc/hosts", &format!("127.0.0.1 localhost\n::1 localhost\n127.0.1.1 {}\n", conf.hostname), 0o644);
+    install_file(
+        "/mnt/etc/hosts",
+        &format!(
+            "127.0.0.1 localhost\n::1 localhost\n127.0.1.1 {}\n",
+            conf.hostname
+        ),
+        0o644,
+    );
 
     // LOCALE
     print_status("Setting locale");

src/install/gpu.rs

@@ -3,13 +3,25 @@ use crate::{config::GPUVendor, pkg::install_pkgs};
 pub fn setup_video_drivers(vendor: &GPUVendor) {
     match vendor {
         GPUVendor::AMD => {
-            install_pkgs(&["xf86-video-amdgpu", "mesa", "lib32-mesa", "vulkan-radeon", "lib32-vulkan-radeon"]);
-        },
+            install_pkgs(&[
+                "xf86-video-amdgpu",
+                "mesa",
+                "lib32-mesa",
+                "vulkan-radeon",
+                "lib32-vulkan-radeon",
+            ]);
+        }
         GPUVendor::NVIDIA => {
             install_pkgs(&["nvidia", "nvidia-utils", "lib32-nvidia-utils"]);
-        },
+        }
         GPUVendor::INTEL => {
-            install_pkgs(&["xf86-video-intel2", "mesa", "lib32-mesa", "vulkan-intel", "lib32-vulkan-intel"]);
-        },
+            install_pkgs(&[
+                "xf86-video-intel2",
+                "mesa",
+                "lib32-mesa",
+                "vulkan-intel",
+                "lib32-vulkan-intel",
+            ]);
+        }
     }
 }

src/install/mod.rs

@@ -1,6 +1,5 @@
 // TODO : Autojoin docker swarm
 // TODO : Autojoin teleport
-// TODO : Firewall
 
 // DRIVE SELECTION
 
@@ -9,6 +8,7 @@ use boot::setup_bootloader;
 use desktop::setup_desktop;
 use docker::setup_docker;
 use drives::{format_drives, mount_drives};
+use firewall::setup_firewall;
 use firmware::{setup_fwupd, setup_microcode};
 use first_boot::{first_boot_values, genfstab};
 use gpu::setup_video_drivers;
@@ -28,7 +28,10 @@ pub mod boot;
 pub mod desktop;
 pub mod docker;
 pub mod drives;
+pub mod firewall;
+pub mod firmware;
 pub mod first_boot;
+pub mod gpu;
 pub mod kernel;
 pub mod navos;
 pub mod ollama;
@@ -38,8 +41,6 @@ pub mod ssh;
 pub mod user;
 pub mod virt;
 pub mod zram;
-pub mod firmware;
-pub mod gpu;
 
 use crate::{
     config::InstallConfig,
@@ -142,6 +143,10 @@ pub fn install(conf: InstallConfig) {
         setup_bluetooth();
     }
 
+    if conf.general.firewall.unwrap_or(true) {
+        setup_firewall(conf.ssh.is_some());
+    }
+
     setup_vm();
     if let Some(gpu_vendor) = &conf.general.gpu_driver {
         setup_video_drivers(gpu_vendor);
@@ -155,5 +160,5 @@ pub fn install(conf: InstallConfig) {
         setup_tpm_unlock(&conf.drive);
     }
 
-    println!("{}", "System install complete".paint(Color::Green));
+    println!("{}", "System install complete".bold().paint(Color::Green));
 }

src/install/ssh.rs

@@ -1,4 +1,9 @@
-use crate::{config::SSHConfig, linux::install_file, pkg::install_pkgs, print_status};
+use crate::{
+    config::SSHConfig,
+    linux::{install_file, systemd_service_enable},
+    pkg::install_pkgs,
+    print_status,
+};
 use std::io::Write;
 
 /// Setup SSH on the system
@@ -34,5 +39,7 @@ pub fn setup_ssh(conf: &Option<SSHConfig>) {
                 writeln!(authorized_keys, "{}\n", key.key).unwrap();
             }
         }
+
+        systemd_service_enable("sshd.service");
     }
 }

src/install/user.rs

@@ -1,5 +1,8 @@
 use crate::{
-    config::UserConfig, linux::{arch_chroot, install_file}, pkg::install_pkgs, print_status
+    config::UserConfig,
+    linux::{arch_chroot, install_file},
+    pkg::install_pkgs,
+    print_status,
 };
 
 pub fn change_passwd(user: &str, pw: &str) {

src/install/virt.rs

@@ -1,7 +1,7 @@
 use crate::{
     config::{InstallConfig, InstallMode},
-    linux::{arch_chroot, systemd_service_enable},
-    pkg::install_pkgs,
+    linux::{arch_chroot, run_command, systemd_service_enable},
+    pkg::install_pkgs, print_status,
 };
 
 pub fn setup_virtualization(conf: &InstallConfig) {
@@ -30,6 +30,19 @@ pub fn setup_virtualization(conf: &InstallConfig) {
     }
 }
 
+/// Setup guest utils if running inside a VM
 pub fn setup_vm() {
-    // TODO : Install guest tools if in virt with  systemd-detect-virt 
+    let is_vm = run_command(&[
+        "systemd-detect-virt", "--vm"
+    ], None, false).0.trim();
+
+    match is_vm {
+        "qemu" | "kvm" => {
+            print_status("Detected KVM. Installing utils");
+            install_pkgs(&["qemu-guest-agent", "spice-vdagent"]);
+            systemd_service_enable("qemu-guest-agent.service");
+            systemd_service_enable("spice-vdagentd.service");
+        },
+        _ => {}
+    }
 }

src/pkg.rs

@@ -3,18 +3,16 @@ use crate::{
     linux::{arch_chroot, run_command},
 };
 
-pub const DESKTOP_PKG: [&str; 16] = [
+pub const DESKTOP_PKG: [&str; 17] = [
     // Desktop
     "plasma",
     "sddm",
-
     // Sound
     "pipewire",
     "pipewire-alsa",
     "pipewire-pulse",
     "pipewire-jack",
     "wireplumber",
-
     // Applications
     "konsole",
     "dolphin",
@@ -24,9 +22,9 @@ pub const DESKTOP_PKG: [&str; 16] = [
     "gwenview",
     "ark",
     "flatpak",
-
     // Misc
-    "navos/navos"
+    "navos/navos",
+    "man",
 ];
 
 pub const SERVER_PKG: [&str; 2] = ["tmux", "navos/navos"];
@@ -60,7 +58,6 @@ pub fn pacstrap(conf: &PackageConfig) {
         "zsh",
         "zsh-completions",
         "zsh-autosuggestions",
-        "man"
     ];
 
     cmd.extend(

src/print.rs

@@ -65,11 +65,20 @@ pub fn print_config(conf: &InstallConfig) {
     }
     if let Some(vendor) = &conf.general.gpu_driver {
         match vendor {
-            crate::config::GPUVendor::AMD => general_info.add_str(format!("🟥 AMD GPU {}", "✔️".paint(Color::Green))),
-            crate::config::GPUVendor::INTEL => general_info.add_str(format!("🟦 Intel GPU {}", "✔️".paint(Color::Green))),
-            crate::config::GPUVendor::NVIDIA => general_info.add_str(format!("🟩 NVIDIA GPU {}", "✔️".paint(Color::Green))),
+            crate::config::GPUVendor::AMD => {
+                general_info.add_str(format!("🟥 AMD GPU {}", "✔️".paint(Color::Green)))
+            }
+            crate::config::GPUVendor::INTEL => {
+                general_info.add_str(format!("🟦 Intel GPU {}", "✔️".paint(Color::Green)))
+            }
+            crate::config::GPUVendor::NVIDIA => {
+                general_info.add_str(format!("🟩 NVIDIA GPU {}", "✔️".paint(Color::Green)))
+            }
         }
     }
+    if conf.general.firewall.unwrap_or(true) {
+        general_info.add_str(format!("🔥 Firewall {}", "✔️".paint(Color::Green)));
+    }
 
     root_info.add_tree("🔨 General", general_info);