From 821eb4bbe1a4fac8d17b4df96ab6441985df58d7 Mon Sep 17 00:00:00 2001
From: David Ostrovsky <david@ostrovsky.org>
Date: Thu, 10 Apr 2014 18:58:09 -0400
Subject: [PATCH] Expose SSH command as plugin extension point

---
 src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java |  338 ++++++++++++++++++++++++++++++++++---------------------
 1 files changed, 207 insertions(+), 131 deletions(-)

diff --git a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
index b6944ea..f7c78d2 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
@@ -22,155 +22,231 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeSet;
 
 import org.apache.sshd.server.Command;
 import org.apache.sshd.server.Environment;
 import org.kohsuke.args4j.Argument;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import com.gitblit.transport.ssh.CommandMetaData;
+import com.gitblit.models.UserModel;
+import com.gitblit.utils.StringUtils;
 import com.gitblit.utils.cli.SubcommandHandler;
 import com.google.common.base.Charsets;
 import com.google.common.base.Strings;
 import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
 
-public class DispatchCommand extends BaseCommand {
+public abstract class DispatchCommand extends BaseCommand {
 
-  @Argument(index = 0, required = false, metaVar = "COMMAND", handler = SubcommandHandler.class)
-  private String commandName;
+	private Logger log = LoggerFactory.getLogger(getClass());
 
-  @Argument(index = 1, multiValued = true, metaVar = "ARG")
-  private List<String> args = new ArrayList<String>();
+	@Argument(index = 0, required = false, metaVar = "COMMAND", handler = SubcommandHandler.class)
+	private String commandName;
 
-  private Set<Class<? extends Command>> commands;
-  private Map<String, Class<? extends Command>> map;
-  private Map<String, Command> root;
+	@Argument(index = 1, multiValued = true, metaVar = "ARG")
+	private List<String> args = new ArrayList<String>();
 
-  public DispatchCommand() {
-	  commands = new HashSet<Class<? extends Command>>();
-  }
+	private final Set<Class<? extends BaseCommand>> commands;
+	private Map<String, Class<? extends BaseCommand>> map;
+	private Map<String, BaseCommand> dispatchers;
 
-  public void registerDispatcher(String name, Command cmd) {
-	  if (root == null) {
-		  root = Maps.newHashMap();
-	  }
-	  root.put(name, cmd);
-  }
-
-  public void registerCommand(Class<? extends Command> cmd) {
-	  if (!cmd.isAnnotationPresent(CommandMetaData.class)) {
-		  throw new RuntimeException(MessageFormat.format("{0} must be annotated with {1}!",
-				  cmd.getName(), CommandMetaData.class.getName()));
-	  }
-	  commands.add(cmd);
-  }
-
-  private Map<String, Class<? extends Command>> getMap() {
-    if (map == null) {
-      map = Maps.newHashMapWithExpectedSize(commands.size());
-      for (Class<? extends Command> cmd : commands) {
-        CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
-        map.put(meta.name(), cmd);
-      }
-    }
-    return map;
-  }
-
-  @Override
-  public void start(Environment env) throws IOException {
-    try {
-      parseCommandLine();
-      if (Strings.isNullOrEmpty(commandName)) {
-        StringWriter msg = new StringWriter();
-        msg.write(usage());
-        throw new UnloggedFailure(1, msg.toString());
-      }
-
-      Command cmd = getCommand();
-      if (cmd instanceof BaseCommand) {
-        BaseCommand bc = (BaseCommand) cmd;
-        if (getName().isEmpty()) {
-          bc.setName(commandName);
-        } else {
-          bc.setName(getName() + " " + commandName);
-        }
-        bc.setArguments(args.toArray(new String[args.size()]));
-      } else if (!args.isEmpty()) {
-        throw new UnloggedFailure(1, commandName + " does not take arguments");
-      }
-
-      provideStateTo(cmd);
-      //atomicCmd.set(cmd);
-      cmd.start(env);
-
-    } catch (UnloggedFailure e) {
-      String msg = e.getMessage();
-      if (!msg.endsWith("\n")) {
-        msg += "\n";
-      }
-      err.write(msg.getBytes(Charsets.UTF_8));
-      err.flush();
-      exit.onExit(e.exitCode);
-    }
-  }
-
-  private Command getCommand() throws UnloggedFailure {
-	if (root != null && root.containsKey(commandName)) {
-		return root.get(commandName);
+	protected DispatchCommand() {
+		commands = new HashSet<Class<? extends BaseCommand>>();
 	}
-	final Class<? extends Command> c = getMap().get(commandName);
-      if (c == null) {
-        String msg =
-            (getName().isEmpty() ? "Gitblit" : getName()) + ": "
-                + commandName + ": not found";
-        throw new UnloggedFailure(1, msg);
-      }
 
-      Command cmd = null;
-      try {
-    	  cmd = c.newInstance();
-      } catch (Exception e) {
-    	  throw new UnloggedFailure(1, MessageFormat.format("Failed to instantiate {0} command", commandName));
-      }
-	return cmd;
-  }
+	protected void registerDispatcher(UserModel user, Class<? extends DispatchCommand> cmd) {
+		if (!cmd.isAnnotationPresent(CommandMetaData.class)) {
+			throw new RuntimeException(MessageFormat.format("{0} must be annotated with {1}!", cmd.getName(),
+					CommandMetaData.class.getName()));
+		}
+		if (dispatchers == null) {
+			dispatchers = Maps.newHashMap();
+		}
 
-  @Override
-protected String usage() {
-    final StringBuilder usage = new StringBuilder();
-    usage.append("Available commands");
-    if (!getName().isEmpty()) {
-      usage.append(" of ");
-      usage.append(getName());
-    }
-    usage.append(" are:\n");
-    usage.append("\n");
+		CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
+		if (meta.admin() && !user.canAdmin()) {
+			log.debug(MessageFormat.format("excluding admin dispatch command {0} for {1}", meta.name(), user.username));
+			return;
+		}
 
-    int maxLength = -1;
-    Map<String, Class<? extends Command>> m = getMap();
-    for (String name : m.keySet()) {
-      maxLength = Math.max(maxLength, name.length());
-    }
-    String format = "%-" + maxLength + "s   %s";
-    for (String name : Sets.newTreeSet(m.keySet())) {
-      final Class<? extends Command> c = m.get(name);
-      usage.append("   ");
-      CommandMetaData meta = c.getAnnotation(CommandMetaData.class);
-      if (meta != null) {
-        usage.append(String.format(format, name,
-            Strings.nullToEmpty(meta.description())));
-      }
-      usage.append("\n");
-    }
-    usage.append("\n");
+		try {
+			DispatchCommand dispatcher = cmd.newInstance();
+			dispatcher.registerCommands(user);
+			dispatchers.put(meta.name(), dispatcher);
+		} catch (Exception e) {
+			log.error("failed to register {} dispatcher", meta.name());
+		}
+	}
 
-    usage.append("See '");
-    if (getName().indexOf(' ') < 0) {
-      usage.append(getName());
-      usage.append(' ');
-    }
-    usage.append("COMMAND --help' for more information.\n");
-    usage.append("\n");
-    return usage.toString();
-  }
+	protected abstract void registerCommands(UserModel user);
+
+
+	/**
+	 * Registers a command as long as the user is permitted to execute it.
+	 *
+	 * @param user
+	 * @param cmd
+	 */
+	protected void registerCommand(UserModel user, Class<? extends BaseCommand> cmd) {
+		if (!cmd.isAnnotationPresent(CommandMetaData.class)) {
+			throw new RuntimeException(MessageFormat.format("{0} must be annotated with {1}!", cmd.getName(),
+					CommandMetaData.class.getName()));
+		}
+		CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
+		if (meta.admin() && !user.canAdmin()) {
+			log.debug(MessageFormat.format("excluding admin command {0} for {1}", meta.name(), user.username));
+			return;
+		}
+		commands.add(cmd);
+	}
+
+	private Map<String, Class<? extends BaseCommand>> getMap() {
+		if (map == null) {
+			map = Maps.newHashMapWithExpectedSize(commands.size());
+			for (Class<? extends BaseCommand> cmd : commands) {
+				CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
+				map.put(meta.name(), cmd);
+			}
+			if (dispatchers != null) {
+				for (Map.Entry<String, BaseCommand> entry : dispatchers.entrySet()) {
+					map.put(entry.getKey(), entry.getValue().getClass());
+				}
+			}
+		}
+		return map;
+	}
+
+	@Override
+	public void start(Environment env) throws IOException {
+		try {
+			parseCommandLine();
+			if (Strings.isNullOrEmpty(commandName)) {
+				StringWriter msg = new StringWriter();
+				msg.write(usage());
+				throw new UnloggedFailure(1, msg.toString());
+			}
+
+			BaseCommand cmd = getCommand();
+			if (getName().isEmpty()) {
+				cmd.setName(commandName);
+			} else {
+				cmd.setName(getName() + " " + commandName);
+			}
+			cmd.setArguments(args.toArray(new String[args.size()]));
+
+			provideStateTo(cmd);
+			// atomicCmd.set(cmd);
+			cmd.start(env);
+
+		} catch (UnloggedFailure e) {
+			String msg = e.getMessage();
+			if (!msg.endsWith("\n")) {
+				msg += "\n";
+			}
+			err.write(msg.getBytes(Charsets.UTF_8));
+			err.flush();
+			exit.onExit(e.exitCode);
+		}
+	}
+
+	private BaseCommand getCommand() throws UnloggedFailure {
+		if (dispatchers != null && dispatchers.containsKey(commandName)) {
+			return dispatchers.get(commandName);
+		}
+		final Class<? extends BaseCommand> c = getMap().get(commandName);
+		if (c == null) {
+			String msg = (getName().isEmpty() ? "Gitblit" : getName()) + ": " + commandName + ": not found";
+			throw new UnloggedFailure(1, msg);
+		}
+
+		BaseCommand cmd = null;
+		try {
+			cmd = c.newInstance();
+		} catch (Exception e) {
+			throw new UnloggedFailure(1, MessageFormat.format("Failed to instantiate {0} command", commandName));
+		}
+		return cmd;
+	}
+
+	@Override
+	public String usage() {
+		Set<String> commands = new TreeSet<String>();
+		Set<String> dispatchers = new TreeSet<String>();
+		int maxLength = -1;
+		Map<String, Class<? extends BaseCommand>> m = getMap();
+		for (String name : m.keySet()) {
+			Class<? extends BaseCommand> c = m.get(name);
+			CommandMetaData meta = c.getAnnotation(CommandMetaData.class);
+			if (meta != null) {
+				if (meta.hidden()) {
+					continue;
+				}
+			}
+
+			maxLength = Math.max(maxLength, name.length());
+			if (DispatchCommand.class.isAssignableFrom(c)) {
+				dispatchers.add(name);
+			} else {
+				commands.add(name);
+			}
+		}
+		String format = "%-" + maxLength + "s   %s";
+
+		final StringBuilder usage = new StringBuilder();
+		if (!commands.isEmpty()) {
+			usage.append("Available commands");
+			if (!getName().isEmpty()) {
+				usage.append(" of ");
+				usage.append(getName());
+			}
+			usage.append(" are:\n");
+			usage.append("\n");
+			for (String name : commands) {
+				final Class<? extends Command> c = m.get(name);
+				CommandMetaData meta = c.getAnnotation(CommandMetaData.class);
+				usage.append("   ");
+				usage.append(String.format(format, name, Strings.nullToEmpty(meta.description())));
+				usage.append("\n");
+			}
+			usage.append("\n");
+		}
+
+		if (!dispatchers.isEmpty()) {
+			usage.append("Available command dispatchers");
+			if (!getName().isEmpty()) {
+				usage.append(" of ");
+				usage.append(getName());
+			}
+			usage.append(" are:\n");
+			usage.append("\n");
+			for (String name : dispatchers) {
+				final Class<? extends BaseCommand> c = m.get(name);
+				CommandMetaData meta = c.getAnnotation(CommandMetaData.class);
+				usage.append("   ");
+				usage.append(String.format(format, name, Strings.nullToEmpty(meta.description())));
+				usage.append("\n");
+			}
+			usage.append("\n");
+		}
+
+		usage.append("See '");
+		if (!StringUtils.isEmpty(getName())) {
+			usage.append(getName());
+			usage.append(' ');
+		}
+		usage.append("COMMAND --help' for more information.\n");
+		usage.append("\n");
+		return usage.toString();
+	}
+
+	protected void provideStateTo(final BaseCommand cmd) {
+		if (cmd instanceof BaseCommand) {
+			cmd.setContext(ctx);
+		}
+		cmd.setInputStream(in);
+		cmd.setOutputStream(out);
+		cmd.setErrorStream(err);
+		cmd.setExitCallback(exit);
+	}
 }

--
Gitblit v1.9.1