James Moger
2014-03-14 8384e02b7d6f1a1ecd6782c6b5f682c109cd6e04
src/main/java/com/gitblit/transport/ssh/FileKeyManager.java
@@ -28,8 +28,8 @@
import com.gitblit.Keys;
import com.gitblit.manager.IRuntimeManager;
import com.gitblit.utils.FileUtils;
import com.google.common.base.Charsets;
import com.google.common.base.Joiner;
import com.google.common.io.Files;
/**
@@ -75,10 +75,8 @@
            return null;
         }
         if (keys.exists()) {
            String str = Files.toString(keys, Charsets.ISO_8859_1);
            String [] entries = str.split("\n");
            List<PublicKey> list = new ArrayList<PublicKey>();
            for (String entry : entries) {
            for (String entry : Files.readLines(keys, Charsets.ISO_8859_1)) {
               if (entry.trim().length() == 0) {
                  // skip blanks
                  continue;
@@ -103,40 +101,87 @@
      return null;
   }
   /**
    * Adds a unique key to the keystore.  This function determines uniqueness
    * by disregarding the comment/description field during key comparisons.
    */
   @Override
   public boolean addKey(String username, String data) {
      try {
         File keys = getKeystore(username);
         Files.append(data + '\n', keys, Charsets.ISO_8859_1);
         String newKey = stripCommentFromKey(data);
         List<String> lines = new ArrayList<String>();
         File keystore = getKeystore(username);
         if (keystore.exists()) {
            for (String entry : Files.readLines(keystore, Charsets.ISO_8859_1)) {
               String line = entry.trim();
               if (line.length() == 0) {
                  // keep blanks
                  lines.add(entry);
                  continue;
               }
               if (line.charAt(0) == '#') {
                  // keep comments
                  lines.add(entry);
                  continue;
               }
               // only add keys that do not match the new key
               String oldKey = stripCommentFromKey(line);
               if (!newKey.equals(oldKey)) {
                  lines.add(entry);
               }
            }
         }
         // add new key
         lines.add(data);
         // write keystore
         String content = Joiner.on("\n").join(lines).trim().concat("\n");
         Files.write(content, keystore, Charsets.ISO_8859_1);
         return true;
      } catch (IOException e) {
         throw new RuntimeException("Cannot add ssh key", e);
      }
   }
   
   /**
    * Removes a key from the keystore.
    */
   @Override
   public boolean removeKey(String username, String data) {
      try {
         String rmKey = stripCommentFromKey(data);
         File keystore = getKeystore(username);
         if (keystore.exists()) {
            String str = Files.toString(keystore, Charsets.ISO_8859_1);
            List<String> keep = new ArrayList<String>();
            String [] entries = str.split("\n");
            for (String entry : entries) {
               if (entry.trim().length() == 0) {
            List<String> lines = new ArrayList<String>();
            for (String entry : Files.readLines(keystore, Charsets.ISO_8859_1)) {
               String line = entry.trim();
               if (line.length() == 0) {
                  // keep blanks
                  keep.add(entry);
                  lines.add(entry);
                  continue;
               }
               if (entry.charAt(0) == '#') {
               if (line.charAt(0) == '#') {
                  // keep comments
                  keep.add(entry);
                  lines.add(entry);
                  continue;
               }
               final String[] parts = entry.split(" ");
               if (!parts[1].equals(data)) {
                  keep.add(entry);
               // only include keys that are NOT rmKey
               String oldKey = stripCommentFromKey(line);
               if (!rmKey.equals(oldKey)) {
                  lines.add(entry);
               }
            }
            if (lines.isEmpty()) {
               keystore.delete();
            } else {
               // write keystore
               String content = Joiner.on("\n").join(lines).trim().concat("\n");
               Files.write(content, keystore, Charsets.ISO_8859_1);
            }
            return true;
         }
@@ -148,7 +193,7 @@
   @Override
   public boolean removeAllKeys(String username) {
      return FileUtils.delete(getKeystore(username));
      return getKeystore(username).delete();
   }
   protected File getKeystore(String username) {
@@ -157,4 +202,11 @@
      File keys = new File(dir, username + ".keys");
      return keys;
   }
   /* Strips the comment from the key data and eliminates whitespace diffs */
   protected String stripCommentFromKey(String data) {
      String [] cols = data.split(" ");
      String key = Joiner.on(" ").join(cols[0], cols[1]);
      return key;
   }
}