* Allow hosts to specify which pubkey to use

* Password-protected pubkeys now supported
* Fix some dangling database pointers

9 files changed, 194 insertions, 35 deletions

diff --git a/res/values/arrays.xml b/res/values/arrays.xml
index 33622f3..82949f4 100644
--- a/res/values/arrays.xml
+++ b/res/values/arrays.xml
@@ -20,5 +20,10 @@
         <item>Use left-side keys</item>
         <item>Don't use</item>
     </string-array>
+    
+    <string-array name="list_pubkeyids">
+    	<item value="-2">Do not use keys</item>
+    	<item value="-1">Use any unlocked key</item>
+    </string-array>
 
 </resources>
diff --git a/res/xml/host_prefs.xml b/res/xml/host_prefs.xml
index ca1c12b..3d1c51b 100644
--- a/res/xml/host_prefs.xml
+++ b/res/xml/host_prefs.xml
@@ -30,12 +30,20 @@
 		android:entries="@array/list_colors"
 		android:entryValues="@array/list_colors"
 		/>
-
+<!--  
 	<CheckBoxPreference
 		android:key="usekeys"
 		android:title="Use SSH keys"
 		/>
+-->
 
+	<ListPreference
+		android:key="pubkeyid"
+		android:title="Use Pubkey Auth"
+		android:entries="@array/list_pubkeyids"
+		android:entryValues="@array/list_pubkeyids"
+		/>
+		
 	<EditTextPreference
 		android:key="postlogin"
 		android:title="Post-login automation"
diff --git a/src/org/connectbot/GeneratePubkeyActivity.java b/src/org/connectbot/GeneratePubkeyActivity.java
index 690c6ec..8115601 100644
--- a/src/org/connectbot/GeneratePubkeyActivity.java
+++ b/src/org/connectbot/GeneratePubkeyActivity.java
@@ -253,14 +253,15 @@ public class GeneratePubkeyActivity extends Activity implements OnEntropyGathere
 				Log.d(TAG, "private: " + PubkeyUtils.formatKey(priv));
 				Log.d(TAG, "public: " + PubkeyUtils.formatKey(pub));
 				
-				PubkeyDatabase hostdb = new PubkeyDatabase(GeneratePubkeyActivity.this);
-				hostdb.createPubkey(null,
+				PubkeyDatabase pubkeydb = new PubkeyDatabase(GeneratePubkeyActivity.this);
+				pubkeydb.createPubkey(null,
 						nickname.getText().toString(),
 						keyType,
 						PubkeyUtils.getEncodedPrivate(priv, secret),
 						PubkeyUtils.getEncodedPublic(pub),
 						encrypted,
-						unlockAtStartup.isChecked());				
+						unlockAtStartup.isChecked());
+				pubkeydb.close();
 			} catch (Exception e) {
 				Log.e(TAG, "Could not generate key pair");
 
diff --git a/src/org/connectbot/HostEditorActivity.java b/src/org/connectbot/HostEditorActivity.java
index d0f628f..0ae921e 100644
--- a/src/org/connectbot/HostEditorActivity.java
+++ b/src/org/connectbot/HostEditorActivity.java
@@ -18,12 +18,14 @@
 
 package org.connectbot;
 
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 
 import org.connectbot.util.HostDatabase;
+import org.connectbot.util.PubkeyDatabase;
 
 import android.content.ContentValues;
 import android.content.Intent;
@@ -33,6 +35,7 @@ import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.os.Bundle;
 import android.preference.CheckBoxPreference;
+import android.preference.ListPreference;
 import android.preference.Preference;
 import android.preference.PreferenceActivity;
 import android.util.Log;
@@ -41,11 +44,12 @@ public class HostEditorActivity extends PreferenceActivity implements OnSharedPr
 
 	
 	public class CursorPreferenceHack implements SharedPreferences {
-
+		
 		protected final String table;
 		protected final int id;
 
 		protected Map<String, String> values = new HashMap<String, String>();
+		protected Map<String, String> pubkeys = new HashMap<String, String>();
 		
 		public CursorPreferenceHack(String table, int id) {
 			this.table = table;
@@ -74,6 +78,21 @@ public class HostEditorActivity extends PreferenceActivity implements OnSharedPr
 			cursor.close();
 			db.close();
 			
+			db = pubkeydb.getReadableDatabase();
+			cursor = db.query(PubkeyDatabase.TABLE_PUBKEYS,
+					new String[] { "_id", PubkeyDatabase.FIELD_PUBKEY_NICKNAME },
+					null, null, null, null, null);
+	
+			if (cursor.moveToFirst()) {
+				do {
+					String pubkeyid = String.valueOf(cursor.getLong(0));
+					String value = cursor.getString(1);
+					pubkeys.put(pubkeyid, value);
+				} while (cursor.moveToNext());
+			}
+			
+			cursor.close();
+			db.close();
 		}
 		
 		public boolean contains(String key) {
@@ -190,6 +209,8 @@ public class HostEditorActivity extends PreferenceActivity implements OnSharedPr
 	}
 	
 	protected HostDatabase hostdb = null;
+	protected PubkeyDatabase pubkeydb = null;
+	
 	protected CursorPreferenceHack pref;
 	
 	@Override
@@ -202,15 +223,25 @@ public class HostEditorActivity extends PreferenceActivity implements OnSharedPr
 		//this.getPreferenceManager().setSharedPreferencesName(uri);
 		
 		this.hostdb = new HostDatabase(this);
+		this.pubkeydb = new PubkeyDatabase(this);
 		
 		this.pref = new CursorPreferenceHack(HostDatabase.TABLE_HOSTS, id);
 		this.pref.registerOnSharedPreferenceChangeListener(this);
 		
 		this.addPreferencesFromResource(R.xml.host_prefs);
 		
-		this.updateSummaries();
-		
+		// Grab all the pubkeys from the database cache we have.
+		ListPreference pubkeyPref = (ListPreference)this.findPreference(HostDatabase.FIELD_HOST_PUBKEYID);
+
+		List<CharSequence> pubkeyNicks = new LinkedList<CharSequence>(Arrays.asList(pubkeyPref.getEntries()));
+		pubkeyNicks.addAll(this.pref.pubkeys.values());
+		pubkeyPref.setEntries((CharSequence[]) pubkeyNicks.toArray(new CharSequence[pubkeyNicks.size()]));
 		
+		List<CharSequence> pubkeyIds = new LinkedList<CharSequence>(Arrays.asList(pubkeyPref.getEntryValues()));
+		pubkeyIds.addAll(this.pref.pubkeys.keySet());
+		pubkeyPref.setEntryValues((CharSequence[]) pubkeyIds.toArray(new CharSequence[pubkeyIds.size()]));
+
+		this.updateSummaries();	
 	}
 	
 	public void onStart() {
@@ -235,6 +266,19 @@ public class HostEditorActivity extends PreferenceActivity implements OnSharedPr
 			Preference pref = this.findPreference(key);
 			if(pref == null) continue;
 			if(pref instanceof CheckBoxPreference) continue;
+			String value = this.pref.getString(key, "");
+			
+			if(key.equals("pubkeyid")) {
+				try {
+					int pubkeyId = Integer.parseInt(value);
+					if (pubkeyId >= 0)
+						pref.setSummary(this.pref.pubkeys.get(this.pref.getString(key, "")));
+					continue;
+				} catch (NumberFormatException nfe) {
+					// Fall through.
+				}
+			}
+			
 			pref.setSummary(this.pref.getString(key, ""));
 		}
 		
diff --git a/src/org/connectbot/HostListActivity.java b/src/org/connectbot/HostListActivity.java
index 0727230..cda2131 100644
--- a/src/org/connectbot/HostListActivity.java
+++ b/src/org/connectbot/HostListActivity.java
@@ -270,7 +270,7 @@ public class HostListActivity extends ListActivity {
 					nickname = String.format("%s@%s:%d", username, hostname, port);
 				}
 				
-				hostdb.createHost(null, nickname, username, hostname, port, HostDatabase.COLOR_GRAY);
+				hostdb.createHost(null, nickname, username, hostname, port, HostDatabase.COLOR_GRAY, HostDatabase.PUBKEYID_ANY);
 				
 				Intent intent = new Intent(HostListActivity.this, ConsoleActivity.class);
 				intent.setData(Uri.parse(String.format("ssh://%s@%s:%s/#%s", username, hostname, port, nickname)));
diff --git a/src/org/connectbot/service/TerminalBridge.java b/src/org/connectbot/service/TerminalBridge.java
index cd1cf0c..2712c49 100644
--- a/src/org/connectbot/service/TerminalBridge.java
+++ b/src/org/connectbot/service/TerminalBridge.java
@@ -21,10 +21,13 @@ package org.connectbot.service;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.security.NoSuchAlgorithmException;
 import java.security.PrivateKey;
 import java.security.PublicKey;
+import java.security.spec.InvalidKeySpecException;
 
 import org.connectbot.TerminalView;
+import org.connectbot.util.HostDatabase;
 import org.connectbot.util.PubkeyDatabase;
 import org.connectbot.util.PubkeyUtils;
 
@@ -269,6 +272,46 @@ public class TerminalBridge implements VDUDisplay, OnKeyListener, InteractiveCal
 		}).start();
 	}
 	
+	/**
+	 * Attempt connection with database row pointed to by cursor.
+	 * @param cursor
+	 * @return true for successful authentication
+	 * @throws NoSuchAlgorithmException
+	 * @throws InvalidKeySpecException
+	 * @throws IOException
+	 */
+	public boolean tryPublicKey(Cursor c) throws NoSuchAlgorithmException, InvalidKeySpecException, IOException {
+		String keyNickname = c.getString(COL_NICKNAME);
+		int encrypted = c.getInt(COL_ENCRYPTED);
+		
+		String password = null;
+		if (encrypted != 0)
+			password = promptHelper.requestStringPrompt(String.format("Password for key '%s'", keyNickname));
+		
+		PrivateKey privKey;
+		try {
+			privKey = PubkeyUtils.decodePrivate(c.getBlob(COL_PRIVATE),
+					c.getString(COL_TYPE), password);
+		} catch (Exception e) {
+			e.printStackTrace();
+			outputLine("Bad password for key '" + keyNickname + "'. Authentication failed.");
+			return false;
+		}
+		
+		PublicKey pubKey = PubkeyUtils.decodePublic(c.getBlob(COL_PUBLIC),
+				c.getString(COL_TYPE));
+		
+		Log.d("TerminalBridge", "Trying key " + PubkeyUtils.formatKey(pubKey));
+		
+		if (connection.authenticateWithPublicKey(username,
+				PubkeyUtils.convertToTrilead(privKey, pubKey)))
+			return true;
+		else
+			outputLine("Authentication method 'publickey' with key " + keyNickname + " failed");
+		
+		return false;
+	}
+	
 	public void handleAuthentication() {
 		try {
 			if (connection.authenticateWithNone(username)) {
@@ -282,31 +325,27 @@ public class TerminalBridge implements VDUDisplay, OnKeyListener, InteractiveCal
 		outputLine("Trying to authenticate");
 		
 		try {
-		
+			long pubkeyId = manager.hostdb.getPubkeyId(nickname);
+			
 			if (!pubkeysExhausted &&
+					pubkeyId != HostDatabase.PUBKEYID_NEVER &&
 					connection.isAuthMethodAvailable(username, AUTH_PUBLICKEY)) {
-				Cursor cursor = pubkeydb.allPubkeys();
-				String keyNickname;
-				PrivateKey privKey;
-				PublicKey pubKey;
-								
-				while (cursor.moveToNext()) {
-					if (cursor.getInt(COL_ENCRYPTED) == 0) {
-						keyNickname = cursor.getString(COL_NICKNAME);
-						privKey = PubkeyUtils.decodePrivate(cursor.getBlob(COL_PRIVATE),
-								cursor.getString(COL_TYPE));
-						pubKey = PubkeyUtils.decodePublic(cursor.getBlob(COL_PUBLIC),
-								cursor.getString(COL_TYPE));
-						
-						Log.d("TerminalBridge", "Trying key " + PubkeyUtils.formatKey(pubKey));
-						
-						if (connection.authenticateWithPublicKey(username,
-								PubkeyUtils.convertToTrilead(privKey, pubKey))) {
-							finishConnection();
-						} else {
-							outputLine("Authentication method 'publickey' with key " + keyNickname + " failed");
+				Cursor cursor;
+				
+				if (pubkeyId == HostDatabase.PUBKEYID_ANY) {
+					cursor = pubkeydb.allPubkeys();
+									
+					while (cursor.moveToNext()) {
+						if (cursor.getInt(COL_ENCRYPTED) == 0) {
+							if (tryPublicKey(cursor))
+								finishConnection();
 						}
 					}
+				} else {
+					cursor = pubkeydb.getPubkey(pubkeyId);
+					if (cursor.moveToFirst())
+						if (tryPublicKey(cursor))
+							finishConnection();
 				}
 				
 				cursor.close();
diff --git a/src/org/connectbot/util/HostDatabase.java b/src/org/connectbot/util/HostDatabase.java
index 586180a..3556875 100644
--- a/src/org/connectbot/util/HostDatabase.java
+++ b/src/org/connectbot/util/HostDatabase.java
@@ -38,7 +38,7 @@ public class HostDatabase extends SQLiteOpenHelper {
 	public final static String TAG = HostDatabase.class.toString();
 	
 	public final static String DB_NAME = "hosts";
-	public final static int DB_VERSION = 10;
+	public final static int DB_VERSION = 11;
 	
 	public final static String TABLE_HOSTS = "hosts";
 	public final static String FIELD_HOST_NICKNAME = "nickname";
@@ -51,12 +51,16 @@ public class HostDatabase extends SQLiteOpenHelper {
 	public final static String FIELD_HOST_COLOR = "color";
 	public final static String FIELD_HOST_USEKEYS = "usekeys";
 	public final static String FIELD_HOST_POSTLOGIN = "postlogin";
+	public final static String FIELD_HOST_PUBKEYID = "pubkeyid";
 
 	public final static String COLOR_RED = "red";
 	public final static String COLOR_GREEN = "green";
 	public final static String COLOR_BLUE = "blue";
 	public final static String COLOR_GRAY = "gray";
 
+	public final static long PUBKEYID_NEVER = -2;
+	public final static long PUBKEYID_ANY = -1;
+	
 	public HostDatabase(Context context) {
 		super(context, DB_NAME, null, DB_VERSION);
 	}
@@ -74,12 +78,13 @@ public class HostDatabase extends SQLiteOpenHelper {
 				+ FIELD_HOST_LASTCONNECT + " INTEGER, "
 				+ FIELD_HOST_COLOR + " TEXT, "
 				+ FIELD_HOST_USEKEYS + " TEXT, "
-				+ FIELD_HOST_POSTLOGIN + ")");
+				+ FIELD_HOST_POSTLOGIN + " TEXT, "
+				+ FIELD_HOST_PUBKEYID + " INTEGER DEFAULT " + PUBKEYID_ANY);
 
 		// insert a few sample hosts, none of which probably connect
 		//this.createHost(db, "connectbot@bravo", "connectbot", "192.168.254.230", 22, COLOR_GRAY);
-		this.createHost(db, "cron@server.example.com", "cron", "server.example.com", 22, COLOR_GRAY);
-		this.createHost(db, "backup@example.net", "backup", "example.net", 22, COLOR_BLUE);
+		this.createHost(db, "cron@server.example.com", "cron", "server.example.com", 22, COLOR_GRAY, PUBKEYID_ANY);
+		this.createHost(db, "backup@example.net", "backup", "example.net", 22, COLOR_BLUE, PUBKEYID_ANY);
 		
 	}
 
@@ -91,6 +96,11 @@ public class HostDatabase extends SQLiteOpenHelper {
 			db.execSQL("DROP TABLE IF EXISTS " + TABLE_HOSTS);
 			onCreate(db);	
 		}
+		
+		if (oldVersion == 10) {
+			db.execSQL("ALTER TABLE " + TABLE_HOSTS
+					+ " ADD COLUMN " + FIELD_HOST_PUBKEYID + " INTEGER DEFAULT " + PUBKEYID_ANY);
+		}
 	}
 	
 	/**
@@ -114,7 +124,7 @@ public class HostDatabase extends SQLiteOpenHelper {
 	 * Create a new host using the given parameters, and return its new
 	 * <code>_id</code> value.
 	 */
-	public long createHost(SQLiteDatabase db, String nickname, String username, String hostname, int port, String color) {
+	public long createHost(SQLiteDatabase db, String nickname, String username, String hostname, int port, String color, long pubkeyId) {
 		// create and insert new host
 		
 		if(db == null) db = this.getWritableDatabase();
@@ -128,6 +138,7 @@ public class HostDatabase extends SQLiteOpenHelper {
 		values.put(FIELD_HOST_USEKEYS, Boolean.toString(true));
 		if(color != null)
 			values.put(FIELD_HOST_COLOR, color);
+		values.put(FIELD_HOST_PUBKEYID, pubkeyId);
 		
 		return db.insert(TABLE_HOSTS, null, values);
 		
@@ -229,5 +240,48 @@ public class HostDatabase extends SQLiteOpenHelper {
 		return known;
 	}
 	
+	/**
+	 * Find the pubkey to use for a given nickname.
+	 */
+	public long getPubkeyId(String nickname) {	
+		long result = PUBKEYID_ANY;
+
+		SQLiteDatabase db = this.getReadableDatabase();
+		Cursor c = db.query(TABLE_HOSTS, new String[] { FIELD_HOST_PUBKEYID },
+				FIELD_HOST_NICKNAME + " = ?", new String[] { nickname }, null, null, null);
+		
+		if (c != null && c.moveToFirst())
+			result = c.getLong(0);
+		
+		return result;
+	}
+	
+	public void setPubkeyId(long hostId, long pubkeyId) {
+		SQLiteDatabase db = this.getWritableDatabase();
+		
+		ContentValues values = new ContentValues();
+		values.put(FIELD_HOST_PUBKEYID, pubkeyId);
+		
+		db.update(TABLE_HOSTS, values, "_id = ?", new String[] { String.valueOf(hostId) });
+		db.close();
+
+		Log.d(TAG, String.format("Updated host id %d to use pubkey id %d", hostId, pubkeyId));
+	}
 	
+	/**
+	 * Unset any hosts using a pubkey that has been deleted.
+	 */
+	public void stopUsingPubkey(long pubkeyId) {
+		if (pubkeyId < 0) return;
+		
+		SQLiteDatabase db = this.getWritableDatabase();
+		
+		ContentValues values = new ContentValues();
+		values.put(FIELD_HOST_PUBKEYID, PUBKEYID_ANY);
+		
+		db.update(TABLE_HOSTS, values, FIELD_HOST_PUBKEYID + " = ?", new String[] { String.valueOf(pubkeyId) });
+		db.close();
+		
+		Log.d(TAG, String.format("Set all hosts using pubkey id %d to -1", pubkeyId));
+	}
 }
diff --git a/src/org/connectbot/util/PubkeyDatabase.java b/src/org/connectbot/util/PubkeyDatabase.java
index 83b2f66..1c3278b 100644
--- a/src/org/connectbot/util/PubkeyDatabase.java
+++ b/src/org/connectbot/util/PubkeyDatabase.java
@@ -53,9 +53,13 @@ public class PubkeyDatabase extends SQLiteOpenHelper {
 	public final static String FIELD_PUBKEY_PUBLIC = "public";
 	public final static String FIELD_PUBKEY_ENCRYPTED = "encrypted";
 	public final static String FIELD_PUBKEY_STARTUP = "startup";
+	
+	private Context context;
 
 	public PubkeyDatabase(Context context) {
 		super(context, DB_NAME, null, DB_VERSION);
+		
+		this.context = context;
 	}
 
 	@Override
@@ -103,6 +107,10 @@ public class PubkeyDatabase extends SQLiteOpenHelper {
 	public void deletePubkey(long id) {
 		SQLiteDatabase db = this.getWritableDatabase();
 		db.delete(TABLE_PUBKEYS, "_id = ?", new String[] { Long.toString(id) });
+		
+		HostDatabase hostdb = new HostDatabase(context);
+		hostdb.stopUsingPubkey(id);
+		hostdb.close();
 	}
 	
 	/**
diff --git a/src/org/connectbot/util/PubkeyUtils.java b/src/org/connectbot/util/PubkeyUtils.java
index b62870d..8702987 100644
--- a/src/org/connectbot/util/PubkeyUtils.java
+++ b/src/org/connectbot/util/PubkeyUtils.java
@@ -112,7 +112,7 @@ public class PubkeyUtils {
 	}
 	
 	public static PrivateKey decodePrivate(byte[] encoded, String keyType, String secret) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidKeySpecException {
-		if (secret.length() > 0)
+		if (secret != null && secret.length() > 0)
 			return decodePrivate(decrypt(encoded, secret), keyType);
 		else
 			return decodePrivate(encoded, keyType);