From af93157a7651f0e593ae3716f57cd3d620d0ce96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E9=80=9A=E6=B4=B2?= <carter.li@eoitek.com>
Date: Fri, 14 Oct 2022 18:38:18 +0800
Subject: [PATCH] TerminalShell: handle nested shell

---
 .../terminalshell/terminalshell_linux.c       |  2 +-
 .../terminalshell/terminalshell_windows.cpp   | 75 +++++++++++++++++--
 2 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/src/detection/terminalshell/terminalshell_linux.c b/src/detection/terminalshell/terminalshell_linux.c
index 0765d1a18..0e80e743a 100644
--- a/src/detection/terminalshell/terminalshell_linux.c
+++ b/src/detection/terminalshell/terminalshell_linux.c
@@ -164,7 +164,7 @@ static void getTerminalFromEnv(FFTerminalShellResult* result)
         ffStrbufIgnCaseCompS(&result->terminalProcessName, "0") != 0
     ) return;
 
-    char* term = NULL;
+    const char* term = NULL;
 
     //SSH
     if(getenv("SSH_CONNECTION") != NULL)
diff --git a/src/detection/terminalshell/terminalshell_windows.cpp b/src/detection/terminalshell/terminalshell_windows.cpp
index 0f72942cd..a6d02c612 100644
--- a/src/detection/terminalshell/terminalshell_windows.cpp
+++ b/src/detection/terminalshell/terminalshell_windows.cpp
@@ -57,7 +57,7 @@ static uint32_t getShellInfo(FFTerminalShellResult* result, uint32_t pid)
 {
     uint32_t ppid;
 
-    if(!getProcessInfo(pid, &ppid, &result->shellProcessName, &result->shellExe))
+    if(pid == 0 || !getProcessInfo(pid, &ppid, &result->shellProcessName, &result->shellExe))
         return 0;
     result->shellExeName = result->shellExe.chars + ffStrbufLastIndexC(&result->shellExe, '\\') + 1;
 
@@ -76,6 +76,8 @@ static uint32_t getShellInfo(FFTerminalShellResult* result, uint32_t pid)
         ffStrbufSetS(&result->shellPrettyName, "Windows PowerShell ISE");
     else if(ffStrbufIgnCaseCompS(&result->shellPrettyName, "cmd") == 0)
         ffStrbufSetS(&result->shellPrettyName, "Command Prompt");
+    else if(ffStrbufIgnCaseCompS(&result->shellPrettyName, "nu") == 0)
+        ffStrbufSetS(&result->shellPrettyName, "nushell");
     else if(ffStrbufIgnCaseCompS(&result->terminalPrettyName, "explorer") == 0)
     {
         ffStrbufSetS(&result->terminalPrettyName, "Windows Explorer"); // Started without shell
@@ -89,7 +91,7 @@ static uint32_t getTerminalInfo(FFTerminalShellResult* result, uint32_t pid)
 {
     uint32_t ppid;
 
-    if(!getProcessInfo(pid, &ppid, &result->terminalProcessName, &result->terminalExe))
+    if(pid == 0 || !getProcessInfo(pid, &ppid, &result->terminalProcessName, &result->terminalExe))
         return 0;
     result->terminalExeName = result->terminalExe.chars + ffStrbufLastIndexC(&result->terminalExe, '\\');
 
@@ -97,6 +99,24 @@ static uint32_t getTerminalInfo(FFTerminalShellResult* result, uint32_t pid)
     if(ffStrbufEndsWithIgnCaseS(&result->terminalPrettyName, ".exe"))
         ffStrbufSubstrBefore(&result->terminalPrettyName, result->terminalPrettyName.length - 4);
 
+    if(
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "pwsh") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "cmd") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "bash") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "zsh") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "fish") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "nu") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "powershell") == 0 ||
+        ffStrbufIgnCaseCompS(&result->terminalPrettyName, "powershell_ise") == 0
+    ) {
+        //We are nested shell
+        ffStrbufClear(&result->terminalProcessName);
+        ffStrbufClear(&result->terminalPrettyName);
+        ffStrbufClear(&result->terminalExe);
+        result->terminalExeName = NULL;
+        return getTerminalInfo(result, ppid);
+    }
+
     if(ffStrbufIgnCaseCompS(&result->terminalPrettyName, "WindowsTerminal") == 0)
         ffStrbufSetS(&result->terminalPrettyName, "Windows Terminal");
     else if(ffStrbufIgnCaseCompS(&result->terminalPrettyName, "conhost") == 0)
@@ -107,6 +127,50 @@ static uint32_t getTerminalInfo(FFTerminalShellResult* result, uint32_t pid)
     return ppid;
 }
 
+static void getTerminalFromEnv(FFTerminalShellResult* result)
+{
+    if(
+        result->terminalProcessName.length > 0 &&
+        !ffStrbufStartsWithIgnCaseS(&result->terminalProcessName, "login") &&
+        ffStrbufIgnCaseCompS(&result->terminalProcessName, "(login)") != 0 &&
+        ffStrbufIgnCaseCompS(&result->terminalProcessName, "systemd") != 0 &&
+        ffStrbufIgnCaseCompS(&result->terminalProcessName, "init") != 0 &&
+        ffStrbufIgnCaseCompS(&result->terminalProcessName, "(init)") != 0 &&
+        ffStrbufIgnCaseCompS(&result->terminalProcessName, "0") != 0
+    ) return;
+
+    const char* term = NULL;
+
+    //SSH
+    if(getenv("SSH_CONNECTION") != NULL)
+        term = getenv("SSH_TTY");
+
+    //Windows Terminal
+    if(!term && (
+        getenv("WT_SESSION") != NULL ||
+        getenv("WT_PROFILE_ID") != NULL
+    )) term = "Windows Terminal";
+
+    //Alacritty
+    if(!term && (
+        getenv("ALACRITTY_SOCKET") != NULL ||
+        getenv("ALACRITTY_LOG") != NULL ||
+        getenv("ALACRITTY_WINDOW_ID") != NULL
+    )) term = "Alacritty";
+
+    //Normal Terminal
+    if(!term)
+        term = getenv("TERM");
+
+    if(term)
+    {
+        ffStrbufSetS(&result->terminalProcessName, term);
+        ffStrbufSetS(&result->terminalPrettyName, term);
+        ffStrbufSetS(&result->terminalExe, term);
+        result->terminalExeName = "";
+    }
+}
+
 const FFTerminalShellResult* ffDetectTerminalShell(const FFinstance* instance)
 {
     FF_UNUSED(instance);
@@ -142,12 +206,9 @@ const FFTerminalShellResult* ffDetectTerminalShell(const FFinstance* instance)
         goto exit;
 
     ppid = getShellInfo(&result, ppid);
-    if(ppid == 0)
-        goto exit;
-
-    // TODO: handle nested shells
-
     getTerminalInfo(&result, ppid);
+    if(result.terminalProcessName.length == 0)
+        getTerminalFromEnv(&result);
 
 exit:
     ffThreadMutexUnlock(&mutex);