From a2eff67d44e2d81b9a435d713faa4c31b0c40b0c Mon Sep 17 00:00:00 2001
From: Lorenzo Migliorero <lorenzo.migliorero@wetransfer.com>
Date: Thu, 25 Jul 2024 22:52:39 +0200
Subject: [PATCH] feat: generate ssh key

---
 .../settings/ssh-keys/add-ssh-key.tsx         | 29 +++++++++++-
 server/api/routers/ssh-key.ts                 |  4 ++
 server/utils/filesystem/ssh.ts                | 46 +++++++++++++------
 3 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/components/dashboard/settings/ssh-keys/add-ssh-key.tsx b/components/dashboard/settings/ssh-keys/add-ssh-key.tsx
index 680873a4..cc6e03ab 100644
--- a/components/dashboard/settings/ssh-keys/add-ssh-key.tsx
+++ b/components/dashboard/settings/ssh-keys/add-ssh-key.tsx
@@ -42,6 +42,8 @@ export const AddSSHKey = ({ children }: Props) => {
 	const { mutateAsync, isError, error, isLoading } =
 		api.sshKey.create.useMutation();
 
+	const generateMutation = api.sshKey.generate.useMutation();
+
 	const form = useForm<SSHKey>({
 		resolver: zodResolver(sshKeyCreate),
 	});
@@ -71,8 +73,31 @@ export const AddSSHKey = ({ children }: Props) => {
 			<DialogContent className="max-h-screen  overflow-y-auto sm:max-w-2xl">
 				<DialogHeader>
 					<DialogTitle>SSH Key</DialogTitle>
-					<DialogDescription>
-						In this section you can add an SSH key
+					<DialogDescription className="space-y-4">
+						<div>
+							In this section you can add one of your keys or generate a new
+							one.
+						</div>
+						<Button
+							variant={"secondary"}
+							isLoading={generateMutation.isLoading}
+							className="max-sm:w-full"
+							onClick={async () => {
+								await generateMutation
+									.mutateAsync()
+									.then(async (data) => {
+										toast.success("SSH Key Generated");
+										form.setValue("privateKey", data.privateKey);
+										form.setValue("publicKey", data.publicKey);
+									})
+									.catch(() => {
+										toast.error("Error to generate the SSH Key");
+									});
+							}}
+							type="button"
+						>
+							Generate SSH Key
+						</Button>
 					</DialogDescription>
 				</DialogHeader>
 				{isError && <AlertBlock type="error">{error?.message}</AlertBlock>}
diff --git a/server/api/routers/ssh-key.ts b/server/api/routers/ssh-key.ts
index 68229f83..d51e31c1 100644
--- a/server/api/routers/ssh-key.ts
+++ b/server/api/routers/ssh-key.ts
@@ -10,6 +10,7 @@ import {
 	apiRemoveSshKey,
 	apiUpdateSshKey,
 } from "@/server/db/schema";
+import { generateSSHKey } from "@/server/utils/filesystem/ssh";
 import { TRPCError } from "@trpc/server";
 import {
 	createSshKey,
@@ -49,6 +50,9 @@ export const sshRouter = createTRPCRouter({
 	all: adminProcedure.query(async () => {
 		return await db.query.sshKeys.findMany({});
 	}),
+	generate: protectedProcedure.mutation(async () => {
+		return await generateSSHKey();
+	}),
 	update: adminProcedure.input(apiUpdateSshKey).mutation(async ({ input }) => {
 		try {
 			return await updateSSHKeyById(input);
diff --git a/server/utils/filesystem/ssh.ts b/server/utils/filesystem/ssh.ts
index 09a3a73d..2d3b63d0 100644
--- a/server/utils/filesystem/ssh.ts
+++ b/server/utils/filesystem/ssh.ts
@@ -3,6 +3,29 @@ import * as path from "node:path";
 import { SSH_PATH } from "@/server/constants";
 import { spawnAsync } from "../process/spawnAsync";
 
+const readSSHKey = async (id: string) => {
+	try {
+		if (!fs.existsSync(SSH_PATH)) {
+			fs.mkdirSync(SSH_PATH, { recursive: true });
+		}
+
+		return {
+			privateKey: fs
+				.readFileSync(path.join(SSH_PATH, `${id}_rsa`), {
+					encoding: "utf-8",
+				})
+				.trim(),
+			publicKey: fs
+				.readFileSync(path.join(SSH_PATH, `${id}_rsa.pub`), {
+					encoding: "utf-8",
+				})
+				.trim(),
+		};
+	} catch (error) {
+		throw error;
+	}
+};
+
 export const saveSSHKey = async (
 	id: string,
 	publicKey: string,
@@ -24,21 +47,23 @@ export const saveSSHKey = async (
 	publicKeyStream.end();
 };
 
-export const generateSSHKey = async (id: string) => {
+export const generateSSHKey = async () => {
 	const applicationDirectory = SSH_PATH;
 
 	if (!fs.existsSync(applicationDirectory)) {
 		fs.mkdirSync(applicationDirectory, { recursive: true });
 	}
 
-	const keyPath = path.join(applicationDirectory, `${id}_rsa`);
+	const keyPath = path.join(applicationDirectory, "temp_rsa");
 
 	if (fs.existsSync(`${keyPath}`)) {
 		fs.unlinkSync(`${keyPath}`);
 	}
+
 	if (fs.existsSync(`${keyPath}.pub`)) {
 		fs.unlinkSync(`${keyPath}.pub`);
 	}
+
 	const args = [
 		"-t",
 		"rsa",
@@ -46,25 +71,18 @@ export const generateSSHKey = async (id: string) => {
 		"4096",
 		"-C",
 		"dokploy",
+		"-m",
+		"PEM",
 		"-f",
 		keyPath,
 		"-N",
 		"",
 	];
+
 	try {
 		await spawnAsync("ssh-keygen", args);
-		return keyPath;
-	} catch (error) {
-		throw error;
-	}
-};
-export const readSSHPublicKey = async (id: string) => {
-	try {
-		if (!fs.existsSync(SSH_PATH)) {
-			fs.mkdirSync(SSH_PATH, { recursive: true });
-		}
-		const keyPath = path.join(SSH_PATH, `${id}_rsa.pub`);
-		const data = fs.readFileSync(keyPath, { encoding: "utf-8" });
+		const data = await readSSHKey("temp");
+		await removeSSHKey("temp");
 		return data;
 	} catch (error) {
 		throw error;