Add SSH TOTP authentication support (#350)

* Add SSH TOTP authentication support - Implement keyboard-interactive authentication for SSH connections - Add TOTP dialog component for Terminal and File Manager - Handle TOTP prompts in WebSocket and HTTP connections - Disable Server Stats for TOTP-enabled servers - Add i18n support for TOTP-related messages * Update src/backend/ssh/server-stats.ts Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update src/backend/ssh/file-manager.ts Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --------- Co-authored-by: ZacharyZcR <zacharyzcr1984@gmail.com> Co-authored-by: Karmaa <88517757+LukeGus@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2025-10-08 09:01:48 +08:00
parent 24d622d8e4
commit fde64bd3df
10 changed files with 511 additions and 16 deletions
--- a/src/ui/Desktop/Apps/File
+++ b/src/ui/Desktop/Apps/File
@@ -14,6 +14,7 @@ import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
 import { toast } from "sonner";
 import { useTranslation } from "react-i18next";
+import { TOTPDialog } from "@/ui/components/TOTPDialog";
 import {
  Upload,
  FolderPlus,
@@ -38,6 +39,7 @@ import {
  renameSSHItem,
  moveSSHItem,
  connectSSH,
+  verifySSHTOTP,
  getSSHStatus,
  keepSSHAlive,
  identifySSHSymlink,
@@ -98,6 +100,9 @@ function FileManagerContent({ initialHost, onClose }: FileManagerProps) {
  const [searchQuery, setSearchQuery] = useState("");
  const [lastRefreshTime, setLastRefreshTime] = useState<number>(0);
  const [viewMode, setViewMode] = useState<"grid" | "list">("grid");
+  const [totpRequired, setTotpRequired] = useState(false);
+  const [totpSessionId, setTotpSessionId] = useState<string | null>(null);
+  const [totpPrompt, setTotpPrompt] = useState<string>("");
  const [pinnedFiles, setPinnedFiles] = useState<Set<string>>(new Set());
  const [sidebarRefreshTrigger, setSidebarRefreshTrigger] = useState(0);
  const [isClosing, setIsClosing] = useState<boolean>(false);
@@ -288,6 +293,14 @@ function FileManagerContent({ initialHost, onClose }: FileManagerProps) {
        userId: currentHost.userId,
      });

+      if (result?.requires_totp) {
+        setTotpRequired(true);
+        setTotpSessionId(sessionId);
+        setTotpPrompt(result.prompt || "Verification code:");
+        setIsLoading(false);
+        return;
+      }
+
      setSshSessionId(sessionId);

      try {
@@ -1238,6 +1251,47 @@ function FileManagerContent({ initialHost, onClose }: FileManagerProps) {
    setEditingFile(null);
  }

+  async function handleTotpSubmit(code: string) {
+    if (!totpSessionId || !code) return;
+
+    try {
+      setIsLoading(true);
+      const result = await verifySSHTOTP(totpSessionId, code);
+
+      if (result?.status === "success") {
+        setTotpRequired(false);
+        setTotpPrompt("");
+        setSshSessionId(totpSessionId);
+        setTotpSessionId(null);
+
+        try {
+          const response = await listSSHFiles(totpSessionId, currentPath);
+          const files = Array.isArray(response)
+            ? response
+            : response?.files || [];
+          setFiles(files);
+          clearSelection();
+          initialLoadDoneRef.current = true;
+          toast.success(t("fileManager.connectedSuccessfully"));
+        } catch (dirError: any) {
+          console.error("Failed to load initial directory:", dirError);
+        }
+      }
+    } catch (error: any) {
+      console.error("TOTP verification failed:", error);
+      toast.error(t("fileManager.totpVerificationFailed"));
+    } finally {
+      setIsLoading(false);
+    }
+  }
+
+  function handleTotpCancel() {
+    setTotpRequired(false);
+    setTotpPrompt("");
+    setTotpSessionId(null);
+    if (onClose) onClose();
+  }
+
  function generateUniqueName(
    baseName: string,
    type: "file" | "directory",
@@ -1806,6 +1860,13 @@ function FileManagerContent({ initialHost, onClose }: FileManagerProps) {
          />
        </div>
      </div>
+
+      <TOTPDialog
+        isOpen={totpRequired}
+        prompt={totpPrompt}
+        onSubmit={handleTotpSubmit}
+        onCancel={handleTotpCancel}
+      />
    </div>
  );
 }
--- a/src/ui/Desktop/Apps/Server/Server.tsx
+++ b/src/ui/Desktop/Apps/Server/Server.tsx
@@ -119,11 +119,16 @@ export function Server({
          setMetrics(data);
          setShowStatsUI(true);
        }
-      } catch (error) {
+      } catch (error: any) {
        if (!cancelled) {
          setMetrics(null);
          setShowStatsUI(false);
-          toast.error(t("serverStats.failedToFetchMetrics"));
+          if (error?.code === "TOTP_REQUIRED" ||
+              (error?.response?.status === 403 && error?.response?.data?.error === "TOTP_REQUIRED")) {
+            toast.error(t("serverStats.totpUnavailable"));
+          } else {
+            toast.error(t("serverStats.failedToFetchMetrics"));
+          }
        }
      } finally {
        if (!cancelled) {
@@ -210,17 +215,28 @@ export function Server({
                    setMetrics(data);
                    setShowStatsUI(true);
                  } catch (error: any) {
-                    if (error?.response?.status === 503) {
+                    if (error?.code === "TOTP_REQUIRED" ||
+                        (error?.response?.status === 403 && error?.response?.data?.error === "TOTP_REQUIRED")) {
+                      toast.error(t("serverStats.totpUnavailable"));
+                      setMetrics(null);
+                      setShowStatsUI(false);
+                    } else if (error?.response?.status === 503 || error?.status === 503) {
                      setServerStatus("offline");
-                    } else if (error?.response?.status === 504) {
+                      setMetrics(null);
+                      setShowStatsUI(false);
+                    } else if (error?.response?.status === 504 || error?.status === 504) {
                      setServerStatus("offline");
-                    } else if (error?.response?.status === 404) {
+                      setMetrics(null);
+                      setShowStatsUI(false);
+                    } else if (error?.response?.status === 404 || error?.status === 404) {
                      setServerStatus("offline");
+                      setMetrics(null);
+                      setShowStatsUI(false);
                    } else {
                      setServerStatus("offline");
+                      setMetrics(null);
+                      setShowStatsUI(false);
                    }
-                    setMetrics(null);
-                    setShowStatsUI(false);
                  } finally {
                    setIsRefreshing(false);
                  }
--- a/src/ui/Desktop/Apps/Terminal/Terminal.tsx
+++ b/src/ui/Desktop/Apps/Terminal/Terminal.tsx
@@ -13,6 +13,7 @@ import { WebLinksAddon } from "@xterm/addon-web-links";
 import { useTranslation } from "react-i18next";
 import { toast } from "sonner";
 import { getCookie, isElectron } from "@/ui/main-axios.ts";
+import { TOTPDialog } from "@/ui/components/TOTPDialog";

 interface SSHTerminalProps {
  hostConfig: any;
@@ -55,6 +56,8 @@ export const Terminal = forwardRef<any, SSHTerminalProps>(function SSHTerminal(
  const [isConnecting, setIsConnecting] = useState(false);
  const [connectionError, setConnectionError] = useState<string | null>(null);
  const [isAuthenticated, setIsAuthenticated] = useState(false);
+  const [totpRequired, setTotpRequired] = useState(false);
+  const [totpPrompt, setTotpPrompt] = useState<string>("");
  const isVisibleRef = useRef<boolean>(false);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const reconnectAttempts = useRef(0);
@@ -102,6 +105,25 @@ export const Terminal = forwardRef<any, SSHTerminalProps>(function SSHTerminal(
    } catch (_) {}
  }

+  function handleTotpSubmit(code: string) {
+    if (webSocketRef.current && code) {
+      webSocketRef.current.send(
+        JSON.stringify({
+          type: "totp_response",
+          data: { code },
+        }),
+      );
+      setTotpRequired(false);
+      setTotpPrompt("");
+    }
+  }
+
+  function handleTotpCancel() {
+    setTotpRequired(false);
+    setTotpPrompt("");
+    if (onClose) onClose();
+  }
+
  function scheduleNotify(cols: number, rows: number) {
    if (!(cols > 0 && rows > 0)) return;
    pendingSizeRef.current = { cols, rows };
@@ -422,6 +444,9 @@ export const Terminal = forwardRef<any, SSHTerminalProps>(function SSHTerminal(
          if (onClose) {
            onClose();
          }
+        } else if (msg.type === "totp_required") {
+          setTotpRequired(true);
+          setTotpPrompt(msg.prompt || "Verification code:");
        }
      } catch (error) {
        toast.error(t("terminal.messageParseError"));
@@ -729,6 +754,13 @@ export const Terminal = forwardRef<any, SSHTerminalProps>(function SSHTerminal(
          </div>
        </div>
      )}
+
+      <TOTPDialog
+        isOpen={totpRequired}
+        prompt={totpPrompt}
+        onSubmit={handleTotpSubmit}
+        onCancel={handleTotpCancel}
+      />
    </div>
  );
 });
--- a/src/ui/components/TOTPDialog.tsx
+++ b/src/ui/components/TOTPDialog.tsx
@@ -0,0 +1,81 @@
+import React from "react";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Label } from "@/components/ui/label";
+import { Shield } from "lucide-react";
+import { useTranslation } from "react-i18next";
+
+interface TOTPDialogProps {
+  isOpen: boolean;
+  prompt: string;
+  onSubmit: (code: string) => void;
+  onCancel: () => void;
+}
+
+export function TOTPDialog({
+  isOpen,
+  prompt,
+  onSubmit,
+  onCancel,
+}: TOTPDialogProps) {
+  const { t } = useTranslation();
+
+  if (!isOpen) return null;
+
+  return (
+    <div className="fixed inset-0 flex items-center justify-center z-50">
+      <div className="absolute inset-0 bg-black/50" />
+      <div className="bg-dark-bg border-2 border-dark-border rounded-lg p-6 max-w-md w-full mx-4 relative z-10">
+        <div className="mb-4 flex items-center gap-2">
+          <Shield className="w-5 h-5 text-primary" />
+          <h3 className="text-lg font-semibold">
+            {t("terminal.totpRequired")}
+          </h3>
+        </div>
+        <p className="text-muted-foreground text-sm mb-4">{prompt}</p>
+        <form
+          onSubmit={(e) => {
+            e.preventDefault();
+            const input = e.currentTarget.elements.namedItem(
+              "totpCode",
+            ) as HTMLInputElement;
+            if (input && input.value.trim()) {
+              onSubmit(input.value.trim());
+            }
+          }}
+          className="space-y-4"
+        >
+          <div>
+            <Label htmlFor="totpCode">
+              {t("terminal.totpCodeLabel")}
+            </Label>
+            <Input
+              id="totpCode"
+              name="totpCode"
+              type="text"
+              autoFocus
+              maxLength={6}
+              pattern="[0-9]*"
+              inputMode="numeric"
+              placeholder="000000"
+              className="text-center text-lg tracking-widest mt-1.5"
+            />
+          </div>
+          <div className="flex gap-2">
+            <Button type="submit" className="flex-1">
+              {t("terminal.totpVerify")}
+            </Button>
+            <Button
+              type="button"
+              variant="outline"
+              onClick={onCancel}
+              className="flex-1"
+            >
+              {t("common.cancel")}
+            </Button>
+          </div>
+        </form>
+      </div>
+    </div>
+  );
+}
--- a/src/ui/main-axios.ts
+++ b/src/ui/main-axios.ts
@@ -526,8 +526,8 @@ function handleApiError(error: unknown, operation: string): never {

  if (axios.isAxiosError(error)) {
    const status = error.response?.status;
-    const message = error.response?.data?.error || error.message;
-    const code = error.response?.data?.code;
+    const message = error.response?.data?.message || error.response?.data?.error || error.message;
+    const code = error.response?.data?.code || error.response?.data?.error;
    const url = error.config?.url;
    const method = error.config?.method?.toUpperCase();

@@ -554,11 +554,15 @@ function handleApiError(error: unknown, operation: string): never {
      throw new ApiError(errorMessage, 401, "AUTH_REQUIRED");
    } else if (status === 403) {
      authLogger.warn(`Access denied: ${method} ${url}`, errorContext);
-      throw new ApiError(
-        "Access denied. You do not have permission to perform this action.",
+      const apiError = new ApiError(
+        code === "TOTP_REQUIRED"
+          ? message
+          : "Access denied. You do not have permission to perform this action.",
        403,
-        "ACCESS_DENIED",
+        code || "ACCESS_DENIED",
      );
+      (apiError as any).response = error.response;
+      throw apiError;
    } else if (status === 404) {
      apiLogger.warn(`Not found: ${method} ${url}`, errorContext);
      throw new ApiError(
@@ -1057,6 +1061,21 @@ export async function disconnectSSH(sessionId: string): Promise<any> {
  }
 }

+export async function verifySSHTOTP(
+  sessionId: string,
+  totpCode: string,
+): Promise<any> {
+  try {
+    const response = await fileManagerApi.post("/ssh/connect-totp", {
+      sessionId,
+      totpCode,
+    });
+    return response.data;
+  } catch (error) {
+    handleApiError(error, "verify SSH TOTP");
+  }
+}
+
 export async function getSSHStatus(
  sessionId: string,
 ): Promise<{ connected: boolean }> {