From e892136c141b934aa81950256b336f10f6cedc6f Mon Sep 17 00:00:00 2001
From: JMARyA <jmarya@hydrar.de>
Date: Sat, 28 Dec 2024 01:34:41 +0100
Subject: [PATCH] Added SSH configuration support

---
 installs/testinstall.toml | 10 ++++++++++
 src/config.rs             | 14 ++++++++++++++
 src/install/mod.rs        |  4 ++++
 src/install/ssh.rs        | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+)
 create mode 100644 src/install/ssh.rs

diff --git a/installs/testinstall.toml b/installs/testinstall.toml
index f227e54..fb996b5 100644
--- a/installs/testinstall.toml
+++ b/installs/testinstall.toml
@@ -44,3 +44,13 @@ password = "testpass"
 
 # Allow user to use doas as root
 doas_root= true
+
+# SSH Configuration
+[ssh]
+# Config file for sshd
+sshd_config = "/etc/ssh/sshd_config"
+
+# Install a SSH key for the user as `authorized_keys`
+[[ssh.key]]
+key = "ssh-rsa ... user@host"
+users = ["testuser", "root"]
diff --git a/src/config.rs b/src/config.rs
index e733334..cdc2190 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -11,6 +11,20 @@ pub struct InstallConfig {
     pub pkg: PackageConfig,
     /// User Configuration
     pub user: Vec<UserConfig>,
+    /// SSH Configuration
+    pub ssh: Option<SSHConfig>,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct SSHConfig {
+    pub sshd_config: Option<String>,
+    pub key: Vec<SSHKey>,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct SSHKey {
+    pub key: String,
+    pub users: Vec<String>,
 }
 
 #[derive(Debug, Deserialize)]
diff --git a/src/install/mod.rs b/src/install/mod.rs
index c64e71c..636059c 100644
--- a/src/install/mod.rs
+++ b/src/install/mod.rs
@@ -11,6 +11,7 @@ use drives::{format_drives, mount_drives};
 use first_boot::{first_boot_values, genfstab};
 use kernel::setup_mkinitcpio;
 use security::{setup_secure_boot, setup_tpm_unlock};
+use ssh::setup_ssh;
 use user::setup_users;
 use yansi::{Color, Paint};
 use zram::setup_zram;
@@ -20,6 +21,7 @@ pub mod drives;
 pub mod first_boot;
 pub mod kernel;
 pub mod security;
+pub mod ssh;
 pub mod user;
 pub mod zram;
 
@@ -65,6 +67,8 @@ pub fn install(conf: InstallConfig) {
     first_boot_values(&conf.general);
     setup_users(&conf.user);
 
+    setup_ssh(&conf.ssh);
+
     setup_bootloader();
 
     match conf.general.mode {
diff --git a/src/install/ssh.rs b/src/install/ssh.rs
new file mode 100644
index 0000000..c19360b
--- /dev/null
+++ b/src/install/ssh.rs
@@ -0,0 +1,33 @@
+use crate::{config::SSHConfig, pkg::install_pkgs};
+use std::io::Write;
+
+pub fn setup_ssh(conf: &Option<SSHConfig>) {
+    if let Some(conf) = conf {
+        install_pkgs(&["openssh"]);
+
+        if let Some(sshd_config) = &conf.sshd_config {
+            let content = std::fs::read_to_string(sshd_config).unwrap();
+            std::fs::write("/mnt/etc/ssh/sshd_config", content).unwrap();
+        }
+
+        for key in &conf.key {
+            for user in &key.users {
+                let path = if user == "root" {
+                    std::fs::create_dir_all("/root/.ssh").unwrap();
+                    "/root/.ssh/authorized_keys".to_string()
+                } else {
+                    std::fs::create_dir_all(&format!("/home/{user}/.ssh")).unwrap();
+                    format!("/home/{user}/.ssh/authorized_keys")
+                };
+
+                let mut authorized_keys = std::fs::OpenOptions::new()
+                    .append(true)
+                    .create(true)
+                    .open(path)
+                    .unwrap();
+
+                writeln!(authorized_keys, "{}", format!("{}\n", key.key)).unwrap();
+            }
+        }
+    }
+}