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