diff --git a/hosts/Haven/default.nix b/hosts/Haven/default.nix
index e68ea6a..5850f4c 100644
--- a/hosts/Haven/default.nix
+++ b/hosts/Haven/default.nix
@@ -40,6 +40,16 @@ in
       tpm2.enable = true;
     };
 
+    # Use Disko to partition and format disks
+    disko = {
+      enable = false;
+      primaryDiskID = "nvme-eui.002538da11443b4d";
+      swapFile = { 
+        enable = true;
+        size = "16G";
+      };
+    };
+
     # Change the default text editor. Options are "emacs", "nano", or "vim".
     editor = "nano";
 
diff --git a/hosts/Shura/default.nix b/hosts/Shura/default.nix
index 3e5e382..ddcf346 100644
--- a/hosts/Shura/default.nix
+++ b/hosts/Shura/default.nix
@@ -53,6 +53,16 @@ in
       tpm2.enable = true;
     };
 
+    # Partition disks using Disko
+    disko = {
+      enable = false;
+      primaryDiskID = "nvme-eui.ace42e0025b66c912ee4ac0000000001";
+      swapFile = { 
+        enable = true;
+        size = "16G";
+      };
+    };
+
     # Change the default text editor. Options are "emacs", "nano", or "vim".
     editor = "nano";
 
diff --git a/hosts/Shura/disko.nix b/modules/system/disko.nix
similarity index 79%
rename from hosts/Shura/disko.nix
rename to modules/system/disko.nix
index 3f8b844..081467b 100644
--- a/hosts/Shura/disko.nix
+++ b/modules/system/disko.nix
@@ -1,6 +1,6 @@
 { lib, config, ... }:
 let
-  cfg = config.disko;
+  cfg = config.aux.system.disko;
 
   standardMountOpts = [
     "compress=zstd"
@@ -9,20 +9,12 @@ let
 in
 {
   options = {
-    disko = {
+    aux.system.disko = {
       enable = lib.mkEnableOption (lib.mdDoc "Enables Disko for disk & partition management.");
-      primaryDisk = lib.mkOption {
-        type = lib.types.attrs;
-        description = "The disk to format using Disko.";
-        default = {
-          name = "nvme0n1";
-          id = "";
-        };
-      };
-      enableTPM = lib.mkOption {
-        type = lib.types.bool;
-        description = "Enables TPM2 support.";
-        default = true;
+      primaryDiskID = lib.mkOption {
+        type = lib.types.str;
+        description = "The ID of the disk to manage using Disko. If possible, use the World Wide Name (WWN), e.g `/dev/disk/by-id/nvme-eui.*`";
+        default = "";
       };
       swapFile = {
         enable = lib.mkEnableOption (lib.mdDoc "Enables the creation of swap files.");
@@ -36,8 +28,15 @@ in
   };
 
   config = lib.mkIf cfg.enable {
+    # Check for blank values
+      assertions = [
+        {
+          assertion = (cfg.primaryDiskID != "");
+          message = "aux.system.disko.primaryDiskID is not set. Please enter a valid disk ID.";
+        }
+      ];
     # Disk management
-    disko.enableConfig = false;
+    disko.enableConfig = true;
     disko.devices = {
       disk = {
         main = {
@@ -66,7 +65,7 @@ in
                   name = "cryptroot";
                   settings = {
                     allowDiscards = true;
-                    crypttabExtraOpts = lib.mkIf cfg.enableTPM [ "tpm2-device=auto" ];
+                    crypttabExtraOpts = lib.mkIf config.aux.system.bootloader.tpm2.enable [ "tpm2-device=auto" ];
                   };
                   content = {
                     type = "btrfs";