nshlib: apply #/$ prompt markers after login

When NSH_PROMPT_STRING_ROOT/USER are empty, keep NSH_PROMPT_STRING at
boot (for example, "nsh> ") so CI/NTFC boot detection still works.
After login, su, or telnet login, replace the last '>' with '#' (euid 0)
or '$' (non-zero euid) and ensure a trailing space.  Refresh readline
after console/telnet login when line editing is enabled.

Signed-off-by: Abhishek Mishra <mishra.abhishek2808@gmail.com>

nshlib/Kconfig

@@ -74,18 +74,20 @@ config NSH_PROMPT_STRING_ROOT
 	default ""
 	depends on SCHED_USER_IDENTITY
 	---help---
-		If non-empty, NSH uses this prompt when the effective UID is zero.
-		If empty, the prompt from NSH_PROMPT_STRING (or ENV/HOSTNAME) is used.
-		Set explicitly for multi-user shells (for example, "nsh# ").
+		Optional full prompt override when the effective UID is zero.
+		If empty, NSH keeps NSH_PROMPT_STRING until login; after login the
+		last '>' in the base prompt becomes '#' (for example, "nsh> "
+		becomes "nsh# "), or '#' is appended when the prompt has no '>'.
 
 config NSH_PROMPT_STRING_USER
 	string "Prompt string for non-root effective UID"
 	default ""
 	depends on SCHED_USER_IDENTITY
 	---help---
-		If non-empty, NSH uses this prompt when the effective UID is non-zero.
-		If empty, the prompt from NSH_PROMPT_STRING (or ENV/HOSTNAME) is used.
-		Set explicitly for multi-user shells (for example, "nsh$ ").
+		Optional full prompt override when the effective UID is non-zero.
+		If empty, NSH keeps NSH_PROMPT_STRING until login; after login the
+		last '>' in the base prompt becomes '$' (for example, "nsh> "
+		becomes "nsh$ "), or '$' is appended when the prompt has no '>'.
 
 config NSH_PROMPT_MAX
 	int "Maximum Size of Prompt String"

nshlib/nsh.h

@@ -822,6 +822,7 @@ int nsh_parse(FAR struct nsh_vtbl_s *vtbl, FAR char *cmdline);
 
 FAR const char *nsh_prompt(void);
 void nsh_update_prompt(void);
+void nsh_update_prompt_after_login(void);
 
 /****************************************************************************
  * Name: nsh_login

nshlib/nsh_identity.c

@@ -377,7 +377,7 @@ int cmd_su(FAR struct nsh_vtbl_s *vtbl, int argc, FAR char **argv)
       return ERROR;
     }
 
-  nsh_update_prompt();
+  nsh_update_prompt_after_login();
   return OK;
 }
 #endif

nshlib/nsh_login.c

@@ -256,7 +256,7 @@ int nsh_login(FAR struct console_stdio_s *pstate)
                   return -1;
                 }
 
-              nsh_update_prompt();
+              nsh_update_prompt_after_login();
 #endif
               return OK;
             }

nshlib/nsh_prompt.c

@@ -31,6 +31,7 @@
 #include <assert.h>
 
 #ifdef CONFIG_SCHED_USER_IDENTITY
+#  include <stdbool.h>
 #  include <unistd.h>
 #endif
 
@@ -49,6 +50,70 @@
 
 static char g_nshprompt[CONFIG_NSH_PROMPT_MAX] = CONFIG_NSH_PROMPT_STRING;
 
+#ifdef CONFIG_SCHED_USER_IDENTITY
+static bool g_nsh_privilege_prompt;
+#endif
+
+/****************************************************************************
+ * Private Functions
+ ****************************************************************************/
+
+#ifdef CONFIG_SCHED_USER_IDENTITY
+
+/****************************************************************************
+ * Name: nsh_apply_privilege_marker
+ *
+ * Description:
+ *   Replace the last '>' in the prompt with the privilege marker ('#' or
+ *   '$').  When no '>' is present, append the marker instead.
+ *
+ ****************************************************************************/
+
+static void nsh_apply_privilege_marker(FAR char *prompt, char marker)
+{
+  size_t len;
+  FAR char *p;
+
+  len = strlen(prompt);
+  for (p = prompt + len; p > prompt; p--)
+    {
+      if (*(p - 1) == '>')
+        {
+          *(p - 1) = marker;
+          return;
+        }
+    }
+
+  if (len + 1 < CONFIG_NSH_PROMPT_MAX)
+    {
+      prompt[len]     = marker;
+      prompt[len + 1] = '\0';
+    }
+}
+
+/****************************************************************************
+ * Name: nsh_ensure_trailing_space
+ *
+ * Description:
+ *   Ensure the prompt ends with a separating space before command input.
+ *
+ ****************************************************************************/
+
+static void nsh_ensure_trailing_space(FAR char *prompt)
+{
+  size_t len;
+
+  len = strlen(prompt);
+  if (len > 0 && prompt[len - 1] != ' ' &&
+      len + 1 < CONFIG_NSH_PROMPT_MAX)
+    {
+      prompt[len]     = ' ';
+      prompt[len + 1] = '\0';
+    }
+}
+
+#endif /* CONFIG_SCHED_USER_IDENTITY */
+
 /****************************************************************************
  * Public Functions
  ****************************************************************************/
@@ -63,9 +128,12 @@ static char g_nshprompt[CONFIG_NSH_PROMPT_MAX] = CONFIG_NSH_PROMPT_STRING;
  *   - non-empty NSH_PROMPT_STRING
  *   - non-empty HOSTNAME and suffix
  *
- *   When SCHED_USER_IDENTITY is enabled and NSH_PROMPT_STRING_ROOT or
- *   NSH_PROMPT_STRING_USER are non-empty, the prompt for the current
- *   effective UID replaces the value from the sources above.
+ *   When SCHED_USER_IDENTITY is enabled, NSH_PROMPT_STRING_ROOT or
+ *   NSH_PROMPT_STRING_USER replace the prompt when non-empty.
+ *
+ *   After login (see nsh_update_prompt_after_login()), when those overrides
+ *   are empty, the last '>' in the prompt is replaced with '#' (euid 0) or
+ *   '$' (non-zero euid), or the marker is appended when no '>' is present.
  *
  * Note that suffix has higher priority when used to help clearly separate
  * prompts from command line inputs.
@@ -102,20 +170,68 @@ void nsh_update_prompt(void)
 #ifdef CONFIG_SCHED_USER_IDENTITY
   if (geteuid() == 0)
     {
+      bool applied = false;
+
+#ifdef CONFIG_NSH_PROMPT_STRING_ROOT
       if (CONFIG_NSH_PROMPT_STRING_ROOT[0] != '\0')
         {
           strlcpy(g_nshprompt, CONFIG_NSH_PROMPT_STRING_ROOT,
                   CONFIG_NSH_PROMPT_MAX);
+          applied = true;
+        }
+
+#endif
+
+      if (!applied && g_nsh_privilege_prompt)
+        {
+          nsh_apply_privilege_marker(g_nshprompt, '#');
         }
     }
-  else if (CONFIG_NSH_PROMPT_STRING_USER[0] != '\0')
+  else
     {
-      strlcpy(g_nshprompt, CONFIG_NSH_PROMPT_STRING_USER,
-              CONFIG_NSH_PROMPT_MAX);
+      bool applied = false;
+
+#ifdef CONFIG_NSH_PROMPT_STRING_USER
+      if (CONFIG_NSH_PROMPT_STRING_USER[0] != '\0')
+        {
+          strlcpy(g_nshprompt, CONFIG_NSH_PROMPT_STRING_USER,
+                  CONFIG_NSH_PROMPT_MAX);
+          applied = true;
+        }
+
+#endif
+
+      if (!applied && g_nsh_privilege_prompt)
+        {
+          nsh_apply_privilege_marker(g_nshprompt, '$');
+        }
+    }
+
+  if (g_nsh_privilege_prompt)
+    {
+      nsh_ensure_trailing_space(g_nshprompt);
     }
 #endif
 }
 
+/****************************************************************************
+ * Name: nsh_update_prompt_after_login
+ *
+ * Description:
+ *   Enable privilege markers in the prompt and refresh it.  Boot and
+ *   no-login sessions keep NSH_PROMPT_STRING (for example, "nsh> ").
+ *
+ ****************************************************************************/
+
+void nsh_update_prompt_after_login(void)
+{
+#ifdef CONFIG_SCHED_USER_IDENTITY
+  g_nsh_privilege_prompt = true;
+#endif
+
+  nsh_update_prompt();
+}
+
 /****************************************************************************
  * Name: nsh_prompt
  *

nshlib/nsh_session.c

@@ -103,6 +103,13 @@ int nsh_session(FAR struct console_stdio_s *pstate,
     }
 #endif /* CONFIG_NSH_TELNET_LOGIN */
 
+#ifdef CONFIG_SCHED_USER_IDENTITY
+  if (login != NSH_LOGIN_NONE)
+    {
+      nsh_update_prompt_after_login();
+    }
+#endif
+
   if (login != NSH_LOGIN_NONE)
     {
       /* Present a greeting and possibly a Message of the Day (MOTD) */

nshlib/nsh_telnetlogin.c

@@ -261,7 +261,7 @@ int nsh_telnetlogin(FAR struct console_stdio_s *pstate)
                   return -1;
                 }
 
-              nsh_update_prompt();
+              nsh_update_prompt_after_login();
 #endif
               return OK;
             }