diff --git a/flake.nix b/flake.nix
index e6c294b..db3c004 100644
--- a/flake.nix
+++ b/flake.nix
@@ -66,7 +66,7 @@
         inherit system;
         specialArgs = { inherit kubenix; };
         modules = [
-          machine.specificConfig
+          machine.nixosModule
           disko.nixosModules.disko
           agenix.nixosModules.default
           ./configuration.nix
diff --git a/machines/default.nix b/machines/default.nix
index 0aa2aab..4d4457f 100644
--- a/machines/default.nix
+++ b/machines/default.nix
@@ -3,51 +3,18 @@
     name = "jefke";
     hostname = "jefke.hyp";
 
-    specificConfig = {
+    nixosModule = {
       custom = {
         dataDisk.enable = true;
+        terraformDatabase.enable = true;
+        k3s.enable = true;
+        disko.osDiskDevice = "/dev/nvme0n1";
 
         ssh = {
           useCertificates = true;
           hostCert = builtins.readFile ./jefke_host_ed25519-cert.pub;
           userCert = builtins.readFile ./jefke_user_ed25519-cert.pub;
         };
-
-        terraformDatabase.enable = true;
-
-        k3s.enable = true;
-      };
-
-      disko.devices = {
-        disk = {
-          vdb = {
-            device = "/dev/nvme0n1";
-            type = "disk";
-            content = {
-              type = "gpt";
-              partitions = {
-                ESP = {
-                  type = "EF00";
-                  size = "500M";
-                  content = {
-                    type = "filesystem";
-                    format = "vfat";
-                    mountpoint = "/boot";
-                  };
-                };
-                root = {
-                  end = "-4G";
-                  content = {
-                    type = "filesystem";
-                    format = "btrfs";
-                    mountpoint = "/";
-                  };
-                };
-                swap = { size = "100%"; };
-              };
-            };
-          };
-        };
       };
     };
   };
@@ -56,46 +23,16 @@
     name = "atlas";
     hostname = "atlas.hyp";
 
-    specificConfig = {
+    nixosModule = {
       custom = {
+        disko.osDiskDevice = "/dev/nvme0n1";
+
         ssh = {
           useCertificates = true;
           hostCert = builtins.readFile ./atlas_host_ed25519-cert.pub;
           userCert = builtins.readFile ./atlas_user_ed25519-cert.pub;
         };
       };
-
-      disko.devices = {
-        disk = {
-          vdb = {
-            device = "/dev/nvme0n1";
-            type = "disk";
-            content = {
-              type = "gpt";
-              partitions = {
-                ESP = {
-                  type = "EF00";
-                  size = "500M";
-                  content = {
-                    type = "filesystem";
-                    format = "vfat";
-                    mountpoint = "/boot";
-                  };
-                };
-                root = {
-                  end = "-4G";
-                  content = {
-                    type = "filesystem";
-                    format = "btrfs";
-                    mountpoint = "/";
-                  };
-                };
-                swap = { size = "100%"; };
-              };
-            };
-          };
-        };
-      };
     };
   };
 }
diff --git a/modules/custom/default.nix b/modules/custom/default.nix
index 29b1813..4cdb9ba 100644
--- a/modules/custom/default.nix
+++ b/modules/custom/default.nix
@@ -1,3 +1,3 @@
 {
-  imports = [ ./terraform-database.nix ./data-disk.nix ./ssh-certificates.nix ./k3s ];
+  imports = [ ./terraform-database.nix ./data-disk.nix ./ssh-certificates.nix ./k3s ./disko.nix ];
 }
diff --git a/modules/custom/disko.nix b/modules/custom/disko.nix
new file mode 100644
index 0000000..4e609c9
--- /dev/null
+++ b/modules/custom/disko.nix
@@ -0,0 +1,43 @@
+{ lib, config, ... }:
+let cfg = config.custom.disko;
+in {
+  options = {
+    custom = {
+      disko.osDiskDevice = lib.mkOption {
+        type = lib.types.str;
+        description = ''
+          The disk device to be used for the operating system.
+        '';
+      };
+    };
+  };
+
+  # TODO: rename this to 'osDisk'. Unfortunately, we would need to run nixos-anywhere again then
+  config.disko.devices.disk.vdb = {
+    device = cfg.osDiskDevice;
+    type = "disk";
+    content = {
+      type = "gpt";
+      partitions = {
+        ESP = {
+          type = "EF00";
+          size = "500M";
+          content = {
+            type = "filesystem";
+            format = "vfat";
+            mountpoint = "/boot";
+          };
+        };
+        root = {
+          end = "-4G";
+          content = {
+            type = "filesystem";
+            format = "btrfs";
+            mountpoint = "/";
+          };
+        };
+        swap = { size = "100%"; };
+      };
+    };
+  };
+}