support for bash like prompt and groups command (#12)

* support for bash like prompt and groups command

* banner update

Author: Vity01

README.md

@@ -24,6 +24,8 @@ There are 3 possible usecases:
 - you can easily add any other HDFS manipulation function
 - there is a command history persisting in history log (~/.hdfs-shell/hdfs-shell.log)
 - **support for relative directory + commands ```cd``` and ```pwd```**
+- advanced commands like ```su```, ```groups```, ```whoami```
+- customizable shell prompt
 
 #### Disadvantages UI against direct calling hdfs dfs function:
 
@@ -72,13 +74,21 @@ For our purposes we also integrated following commands:
 - ```cd```, ```pwd```
 - ```su <username>``` - ***experimental*** -  changes current active user - it won't probably work on secured HDFS (KERBEROS)
 - ```whoami``` - prints effective username
+- ```groups <username1 <username2,...>>``` - eg.```groups hdfs``` prints groups for given users, same as ```hdfs groups my_user my_user2``` functionality
 - ```edit 'my file'``` - see the config below
 
 
 ###### Edit Command
 Since the version 1.0.4 the simple command 'edit' is available. The command gets selected file from HDFS to the local temporary directory and launches the editor. Once the editor saves the file (with a result code 0), the file is uploaded back into HDFS (target file is overwritten).
 By default the editor path is taken from ```$EDITOR``` environment variable. If ```$EDITOR``` is not set, ```vim``` (Linux, Mac) or ```notepad.exe``` (Windows) is used.
 
+###### How to change command (shell) prompt
+HDFS Shell supports customized bash-like prompt setting!
+I implemented support for these switches listed in this [table](https://bash.cyberciti.biz/guide/Changing_bash_prompt) (include colors!, exclude ```\!, \#```).
+You can also use this [online prompt generator](http://ezprompt.net/) to create prompt value of your wish.
+To setup your favorite prompt simply add ```export HDFS_SHELL_PROMPT="value"``` to your .bashrc (or set env variable on Windows) and that's it. Restart HDFS Shell to apply change.
+Default value is currently set to ```\e[36m\u@\h \e[0;39m\e[33m\w\e[0;39m\e[36m\\$ \e[37;0;39m```.
+
 ### Running Daemon mode
 ![Image of HDFS-Shell](https://github.com/avast/hdfs-shell/blob/master/web/screenshot2.png)
 
@@ -101,7 +111,7 @@ For developing, add to JVM args in your IDE launch config dialog:
 
 #### Known limitations & problems
 
-- There is a problem with a parsing of commands containing a file or directory including a space - eg. it's not possible to create directory ```My dir``` using command ```mkdir "My dir"``` . This would be probably resolved with an upgrade to Spring Shell 2. 
+- There is a problem with a parsing of commands containing a file or directory including a space - eg. it's not possible to create directory ```My dir``` using command ```mkdir "My dir"``` . This should be probably resolved with an upgrade to Spring Shell 2.
 - It's not possible to remove root directory (```rm -R dir```) from root (```/```) directory. You have to use absolut path instead (```rm -R /dir```). It's caused by bug in Hadoop. See [HADOOP-15233](https://issues.apache.org/jira/browse/HADOOP-15233) for more details. Removing directory from another cwd is not affected. 
 
 ### Contact

gradle.properties

@@ -1,3 +1,3 @@
 springShellVersion = 1.2.0.RELEASE
 theGroup=com.avast.server
-theVersion=1.0.6
+theVersion=1.0.7

src/main/java/com/avast/server/hdfsshell/commands/ContextCommands.java

@@ -4,20 +4,28 @@
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hdfs.HdfsConfiguration;
+import org.apache.hadoop.hdfs.NameNodeProxies;
 import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.tools.GetUserMappingsProtocol;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.shell.core.CommandMarker;
 import org.springframework.shell.core.annotation.CliAvailabilityIndicator;
 import org.springframework.shell.core.annotation.CliCommand;
 import org.springframework.shell.core.annotation.CliOption;
 import org.springframework.stereotype.Component;
 import org.springframework.util.StringUtils;
 
+import javax.annotation.PostConstruct;
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.stream.Collectors;
 
 @SuppressWarnings("SameParameterValue")
 @Component
 public class ContextCommands implements CommandMarker {
+    private static final Logger logger = LoggerFactory.getLogger(ContextCommands.class);
 
     private String currentDir;
     private Configuration configuration;
@@ -26,13 +34,55 @@ public class ContextCommands implements CommandMarker {
     private boolean showResultCode = false;
     private boolean failOnError;
 
+    private GetUserMappingsProtocol userMappingsProtocol;
 
-    @CliAvailabilityIndicator({"pwd", "cd"})
+    @PostConstruct
+    public void init() {
+        try {
+            final HdfsConfiguration conf = new HdfsConfiguration();
+            userMappingsProtocol = NameNodeProxies.createProxy(conf, FileSystem.getDefaultUri(conf),
+                    GetUserMappingsProtocol.class).getProxy();
+        } catch (Exception e) {
+            logger.error("Failed to create proxy to get user groups", e);
+        }
+    }
+
+    public String[] getGroupsForUser(String username) {
+        if (userMappingsProtocol != null) {
+            try {
+                return userMappingsProtocol.getGroupsForUser(username);
+            } catch (IOException e) {
+                return new String[0];
+            }
+        }
+        return new String[0];
+    }
+
+    @CliAvailabilityIndicator({"pwd", "cd", "groups"})
     public boolean isSimpleAvailable() {
         //always available
         return true;
     }
 
+    @CliCommand(value = "groups", help = "Get groups for user")
+    public String groups(@CliOption(key = {""}) String username) throws IOException {
+        if (StringUtils.isEmpty(username)) {
+            username = whoami();
+        }
+        final StringBuilder result = new StringBuilder();
+        Arrays.stream(username.split("\\W+")).forEach((user) -> {
+                    if (!user.trim().isEmpty()) {
+                        user = user.trim();
+                        if (result.length() > 0) {
+                            result.append(System.lineSeparator());
+                        }
+                        result.append(user).append(" : ").append(Arrays.stream(this.getGroupsForUser(user)).collect(Collectors.joining(" ")));
+                    }
+                }
+        );
+        return result.toString();
+    }
+
     @CliCommand(value = "set", help = "Set switch value")
     public String set(@CliOption(key = {""}, help = "showResultCodeON/showResultCodeOFF") String commandSwitch) {
         if (commandSwitch == null) {

src/main/java/com/avast/server/hdfsshell/ui/ShellBannerProvider.java

@@ -15,18 +15,24 @@
  */
 package com.avast.server.hdfsshell.ui;
 
+import org.springframework.boot.ansi.AnsiColor;
+import org.springframework.boot.ansi.AnsiOutput;
 import org.springframework.core.Ordered;
 import org.springframework.core.annotation.Order;
 import org.springframework.shell.plugin.support.DefaultBannerProvider;
 import org.springframework.stereotype.Component;
 
+import javax.annotation.PostConstruct;
+
 /**
  * @author Jarred Li
  */
 @Component
 @Order(Ordered.HIGHEST_PRECEDENCE)
 public class ShellBannerProvider extends DefaultBannerProvider {
 
+    private SimpleBashPromptInterpreter bashPromptInterpreter;
+
     public String getBanner() {
         return "";//we are using Spring Boot for that
     }
@@ -36,13 +42,18 @@ public String getVersion() {
         return ShellBannerProvider.versionInfo();
     }
 
+    @PostConstruct
+    public void init() {
+        bashPromptInterpreter = new SimpleBashPromptInterpreter.Builder("HDFS-shell CLI \\h").setAddResetEnd(false).build();
+    }
+
 	public String getWelcomeMessage() {
-		return "Welcome to HDFS-shell CLI";
+		return AnsiOutput.toString(AnsiColor.BRIGHT_GREEN, "Welcome to HDFS-shell CLI ", AnsiColor.DEFAULT);
 	}
-	
+
 	@Override
 	public String getProviderName() {
-		return "HDFS-shell CLI";
+		return bashPromptInterpreter.interpret();
 	}
 
     public static String versionInfo() {

src/main/java/com/avast/server/hdfsshell/ui/ShellPromptProvider.java

@@ -1,12 +1,12 @@
 /*
  * Copyright 2011-2012 the original author or authors.
- * 
+ *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  *      http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -16,42 +16,106 @@
 package com.avast.server.hdfsshell.ui;
 
 import com.avast.server.hdfsshell.commands.ContextCommands;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.boot.ansi.AnsiColor;
-import org.springframework.boot.ansi.AnsiOutput;
 import org.springframework.core.Ordered;
 import org.springframework.core.annotation.Order;
 import org.springframework.shell.core.AbstractShell;
 import org.springframework.shell.plugin.support.DefaultPromptProvider;
 import org.springframework.stereotype.Component;
 
+import javax.annotation.PostConstruct;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Locale;
+import java.util.Set;
+import java.util.stream.Collectors;
+
 /**
- * @author Jarred Li
+ * @author Vitasek L.
  */
 @Component
 @Order(Ordered.HIGHEST_PRECEDENCE)
 public class ShellPromptProvider extends DefaultPromptProvider {
 
+    private static final Logger logger = LoggerFactory.getLogger(ShellPromptProvider.class);
+
+    private static final String DEFAULT_PROMPT = "\033[36m\\u@\\h \033[0;39m\033[33m\\w\033[0;39m\033[36m\\$ \033[37;0;39m";
     final ContextCommands contextCommands;
 
+    private SimpleBashPromptInterpreter simpleBashPromptInterpreter;
+
     @Autowired
     public ShellPromptProvider(ContextCommands contextCommands) {
         this.contextCommands = contextCommands;
     }
 
+    @PostConstruct
+    public void init() {
+        setPs1(System.getenv().getOrDefault("HDFS_SHELL_PROMPT", DEFAULT_PROMPT));
+    }
+
+    public void setPs1(String ps1) {
+        simpleBashPromptInterpreter = new SimpleBashPromptInterpreter.Builder(ps1).
+                setAppName(() -> "hdfs-shell").setAppVersion(ShellBannerProvider::versionInfo).
+                setLocale(Locale.ENGLISH).
+                setUsername(this::getWhoami).
+                setIsRoot(this::isRootPrompt).
+                setCwdAbsolut(contextCommands::getCurrentDir).
+                setCwdShort(this::getShortCwd).
+                build();
+    }
+
+    private boolean isRootPrompt() {
+        final String whoami = this.getWhoami();
+        final String[] groupsForUser = contextCommands.getGroupsForUser(whoami);
+        if (groupsForUser.length == 0) { //make guess
+            return "root".equals(whoami) || "hdfs".equals(whoami);
+        }
+        final String[] groups = contextCommands.getConfiguration().get("dfs.permissions.superusergroup", "supergroup").split(",");
+        final Set<String> adminGroups = Arrays.stream(groups).map(String::trim).collect(Collectors.toSet());
+        adminGroups.add("Administrators");//for Windows
+        adminGroups.add("hdfs");//special cases
+        adminGroups.add("root");
+        return Arrays.stream(groupsForUser).anyMatch(adminGroups::contains);
+    }
 
     @Override
     public String getPrompt() {
-        return AbstractShell.shellPrompt =
-                AnsiOutput.toString(AnsiColor.CYAN, "hdfs-shell ") +
-                        AnsiOutput.toString(AnsiColor.YELLOW, contextCommands.getCurrentDir()) +
-                        AnsiOutput.toString(AnsiColor.CYAN, " >", AnsiColor.WHITE);
+        return AbstractShell.shellPrompt = simpleBashPromptInterpreter.interpret();
+    }
+
+    private String getWhoami() {
+        try {
+            return contextCommands.whoami();
+        } catch (IOException e) {
+            logger.error("Failed to get active user", e);
+            return "hdfs-shell";
+        }
     }
 
+//    private String defaultPrompt() {
+//        return AnsiOutput.toString(AnsiColor.CYAN, "hdfs-shell ") +
+//                AnsiOutput.toString(AnsiColor.YELLOW, contextCommands.getCurrentDir()) +
+//                AnsiOutput.toString(AnsiColor.CYAN, " >", AnsiColor.WHITE);
+//    }
+
 
     @Override
     public String getProviderName() {
         return "Hdfs-shell prompt";
     }
 
+    private String getShortCwd() {
+        String currentDir = contextCommands.getCurrentDir();
+        if (currentDir.startsWith("/user/")) {
+            final String userHome = "/user/" + this.getWhoami();//call getWhoami later
+            if (currentDir.startsWith(userHome)) {
+                currentDir = StringUtils.replaceOnce(currentDir, userHome, "~");
+            }
+        }
+        return currentDir;
+    }
 }