Add confirm-use and lifetime constraints to individual pubkeys

git-svn-id: https://connectbot.googlecode.com/svn/trunk/connectbot@334 df292f66-193f-0410-a5fc-6d59da041ff2

7 files changed, 139 insertions, 27 deletions

diff --git a/src/com/trilead/ssh2/AuthAgentCallback.java b/src/com/trilead/ssh2/AuthAgentCallback.java
index ffc3020..c395198 100644
--- a/src/com/trilead/ssh2/AuthAgentCallback.java
+++ b/src/com/trilead/ssh2/AuthAgentCallback.java
@@ -20,9 +20,11 @@ public interface AuthAgentCallback {
 	 *            containing a DSA or RSA private key of
 	 *            the user in Trilead object format.
 	 * @param comment comment associated with this key
+	 * @param confirmUse whether to prompt before using this key
+	 * @param lifetime lifetime in seconds for key to be remembered
 	 * @return success or failure
 	 */
-	boolean addIdentity(Object key, String comment);
+	boolean addIdentity(Object key, String comment, boolean confirmUse, int lifetime);
 
 	/**
 	 * @param publicKey byte blob containing the OpenSSH-format encoded public key
diff --git a/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java b/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java
index 1c6be84..b4ab108 100644
--- a/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java
+++ b/src/com/trilead/ssh2/channel/AuthAgentForwardThread.java
@@ -129,7 +129,10 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 						sendIdentities();
 						break;
 					case SSH2_AGENTC_ADD_IDENTITY:
-						addIdentity(tr);
+						addIdentity(tr, false);
+						break;
+					case SSH2_AGENTC_ADD_ID_CONSTRAINED:
+						addIdentity(tr, true);
 						break;
 					case SSH2_AGENTC_REMOVE_IDENTITY:
 						removeIdentity(tr);
@@ -240,7 +243,7 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 	/**
 	 * @param tr
 	 */
-	private void addIdentity(TypesReader tr) {
+	private void addIdentity(TypesReader tr, boolean checkConstraints) {
 		try
 		{
 			if (failWhenLocked())
@@ -275,7 +278,25 @@ public class AuthAgentForwardThread extends Thread implements IChannelWorkerThre
 				return;
 			}
 
-			if (authAgent.addIdentity(key, comment))
+			boolean confirmUse = false;
+			int lifetime = 0;
+
+			if (checkConstraints) {
+				while (tr.remain() > 0) {
+					int constraint = tr.readByte();
+					if (constraint == SSH_AGENT_CONSTRAIN_CONFIRM)
+						confirmUse = true;
+					else if (constraint == SSH_AGENT_CONSTRAIN_LIFETIME)
+						lifetime = tr.readUINT32();
+					else {
+						// Unknown constraint. Bail.
+						os.write(SSH_AGENT_FAILURE);
+						return;
+					}
+				}
+			}
+
+			if (authAgent.addIdentity(key, comment, confirmUse, lifetime))
 				os.write(SSH_AGENT_SUCCESS);
 			else
 				os.write(SSH_AGENT_FAILURE);
diff --git a/src/org/connectbot/PubkeyListActivity.java b/src/org/connectbot/PubkeyListActivity.java
index 4bcd41b..452e3b4 100644
--- a/src/org/connectbot/PubkeyListActivity.java
+++ b/src/org/connectbot/PubkeyListActivity.java
@@ -90,6 +90,7 @@ public class PubkeyListActivity extends ListActivity implements EventListener {
 	protected TerminalManager bound = null;
 
 	private MenuItem onstartToggle = null;
+	private MenuItem confirmUse = null;
 
 	private ServiceConnection connection = new ServiceConnection() {
 		public void onServiceConnected(ComponentName className, IBinder service) {
@@ -345,7 +346,7 @@ public class PubkeyListActivity extends ListActivity implements EventListener {
 
 		// save this key in-memory if option enabled
 		if(bound.isSavingKeys()) {
-			bound.addKey(pubkey.getNickname(), trileadKey);
+			bound.addKey(pubkey, trileadKey);
 		}
 
 		updateHandler.sendEmptyMessage(-1);
@@ -480,6 +481,19 @@ public class PubkeyListActivity extends ListActivity implements EventListener {
 			}
 		});
 
+		confirmUse = menu.add(R.string.pubkey_confirm_use);
+		confirmUse.setCheckable(true);
+		confirmUse.setChecked(pubkey.isConfirmUse());
+		confirmUse.setOnMenuItemClickListener(new OnMenuItemClickListener() {
+			public boolean onMenuItemClick(MenuItem item) {
+				// toggle confirm use
+				pubkey.setConfirmUse(!pubkey.isConfirmUse());
+				pubkeydb.savePubkey(pubkey);
+				updateHandler.sendEmptyMessage(-1);
+				return true;
+			}
+		});
+
 		MenuItem delete = menu.add(R.string.pubkey_delete);
 		delete.setOnMenuItemClickListener(new OnMenuItemClickListener() {
 			public boolean onMenuItemClick(MenuItem item) {
diff --git a/src/org/connectbot/bean/PubkeyBean.java b/src/org/connectbot/bean/PubkeyBean.java
index 522369d..2db0911 100644
--- a/src/org/connectbot/bean/PubkeyBean.java
+++ b/src/org/connectbot/bean/PubkeyBean.java
@@ -39,11 +39,14 @@ public class PubkeyBean extends AbstractBean {
 	private byte[] publicKey;
 	private boolean encrypted = false;
 	private boolean startup = false;
+	private boolean confirmUse = false;
+	private int lifetime = 0;
 
 	/* Transient values */
 	private boolean unlocked = false;
 	private Object unlockedPrivate = null;
 
+	@Override
 	public String getBeanName() {
 		return BEAN_NAME;
 	}
@@ -116,6 +119,22 @@ public class PubkeyBean extends AbstractBean {
 		return startup;
 	}
 
+	public void setConfirmUse(boolean confirmUse) {
+		this.confirmUse = confirmUse;
+	}
+
+	public boolean isConfirmUse() {
+		return confirmUse;
+	}
+
+	public void setLifetime(int lifetime) {
+		this.lifetime = lifetime;
+	}
+
+	public int getLifetime() {
+		return lifetime;
+	}
+
 	public void setUnlocked(boolean unlocked) {
 		this.unlocked = unlocked;
 	}
@@ -145,6 +164,8 @@ public class PubkeyBean extends AbstractBean {
 		values.put(PubkeyDatabase.FIELD_PUBKEY_PUBLIC, publicKey);
 		values.put(PubkeyDatabase.FIELD_PUBKEY_ENCRYPTED, encrypted ? 1 : 0);
 		values.put(PubkeyDatabase.FIELD_PUBKEY_STARTUP, startup ? 1 : 0);
+		values.put(PubkeyDatabase.FIELD_PUBKEY_CONFIRMUSE, confirmUse ? 1 : 0);
+		values.put(PubkeyDatabase.FIELD_PUBKEY_LIFETIME, lifetime);
 
 		return values;
 	}
diff --git a/src/org/connectbot/service/TerminalManager.java b/src/org/connectbot/service/TerminalManager.java
index 1b0955a..b272845 100644
--- a/src/org/connectbot/service/TerminalManager.java
+++ b/src/org/connectbot/service/TerminalManager.java
@@ -101,6 +101,8 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 
 	private MediaPlayer mediaPlayer;
 
+	private Timer pubkeyTimer;
+
 	private Timer idleTimer;
 	private final long IDLE_TIMEOUT = 300000; // 5 minutes
 
@@ -127,6 +129,8 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 
 		res = getResources();
 
+		pubkeyTimer = new Timer("pubkeyTimer", true);
+
 		hostdb = new HostDatabase(this);
 		pubkeydb = new PubkeyDatabase(this);
 
@@ -139,7 +143,7 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 				PublicKey pubKey = PubkeyUtils.decodePublic(pubkey.getPublicKey(), pubkey.getType());
 				Object trileadKey = PubkeyUtils.convertToTrilead(privKey, pubKey);
 
-				addKey(pubkey.getNickname(), trileadKey);
+				addKey(pubkey, trileadKey);
 			} catch (Exception e) {
 				Log.d(TAG, String.format("Problem adding key '%s' to in-memory cache", pubkey.getNickname()), e);
 			}
@@ -184,6 +188,8 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 		synchronized (this) {
 			if (idleTimer != null)
 				idleTimer.cancel();
+			if (pubkeyTimer != null)
+				pubkeyTimer.cancel();
 		}
 
 		if (wifilock != null && wifilock.isHeld())
@@ -316,18 +322,30 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 		return loadedKeypairs.containsKey(nickname);
 	}
 
-	public void addKey(String nickname, Object trileadKey) {
-		removeKey(nickname);
+	public void addKey(PubkeyBean pubkey, Object trileadKey) {
+		removeKey(pubkey.getNickname());
 
 		byte[] sshPubKey = PubkeyUtils.extractOpenSSHPublic(trileadKey);
 
 		KeyHolder keyHolder = new KeyHolder();
+		keyHolder.bean = pubkey;
 		keyHolder.trileadKey = trileadKey;
 		keyHolder.openSSHPubkey = sshPubKey;
 
-		loadedKeypairs.put(nickname, keyHolder);
+		loadedKeypairs.put(pubkey.getNickname(), keyHolder);
+
+		if (pubkey.getLifetime() > 0) {
+			final String nickname = pubkey.getNickname();
+			pubkeyTimer.schedule(new TimerTask() {
+				@Override
+				public void run() {
+					Log.d(TAG, "Unloading from memory key: " + nickname);
+					removeKey(nickname);
+				}
+			}, pubkey.getLifetime() * 1000);
+		}
 
-		Log.d(TAG, String.format("Added key '%s' to in-memory cache", nickname));
+		Log.d(TAG, String.format("Added key '%s' to in-memory cache", pubkey.getNickname()));
 	}
 
 	public boolean removeKey(String nickname) {
@@ -382,7 +400,7 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 		if (loadedKeypairs.size() > 0) {
 			synchronized (this) {
 				if (idleTimer == null)
-					idleTimer = new Timer(true);
+					idleTimer = new Timer("idleTimer", true);
 
 				idleTimer.schedule(new IdleTask(), IDLE_TIMEOUT);
 			}
@@ -584,6 +602,7 @@ public class TerminalManager extends Service implements BridgeDisconnectedListen
 	}
 
 	public class KeyHolder {
+		public PubkeyBean bean;
 		public Object trileadKey;
 		public byte[] openSSHPubkey;
 	}
diff --git a/src/org/connectbot/transport/SSH.java b/src/org/connectbot/transport/SSH.java
index a54e091..3e3af20 100644
--- a/src/org/connectbot/transport/SSH.java
+++ b/src/org/connectbot/transport/SSH.java
@@ -228,9 +228,13 @@ public class SSH extends AbsTransport implements ConnectionMonitor, InteractiveC
 					// try each of the in-memory keys
 					bridge.outputLine(manager.res
 							.getString(R.string.terminal_auth_pubkey_any));
-					for(String nickname : manager.loadedKeypairs.keySet()) {
-						Object trileadKey = manager.loadedKeypairs.get(nickname).trileadKey;
-						if(this.tryPublicKey(host.getUsername(), nickname, trileadKey)) {
+					for (Entry<String, KeyHolder> entry : manager.loadedKeypairs.entrySet()) {
+						if (entry.getValue().bean.isConfirmUse()
+								&& !promptForPubkeyUse(entry.getKey()))
+							continue;
+
+						if (this.tryPublicKey(host.getUsername(), entry.getKey(),
+								entry.getValue().trileadKey)) {
 							finishConnection();
 							break;
 						}
@@ -291,8 +295,13 @@ public class SSH extends AbsTransport implements ConnectionMonitor, InteractiveC
 		if(manager.isKeyLoaded(pubkey.getNickname())) {
 			// load this key from memory if its already there
 			Log.d(TAG, String.format("Found unlocked key '%s' already in-memory", pubkey.getNickname()));
-			trileadKey = manager.getKey(pubkey.getNickname());
 
+			if (pubkey.isConfirmUse()) {
+				if (promptForPubkeyUse(pubkey.getNickname()))
+					return false;
+			}
+
+			trileadKey = manager.getKey(pubkey.getNickname());
 		} else {
 			// otherwise load key from database and prompt for password as needed
 			String password = null;
@@ -333,7 +342,7 @@ public class SSH extends AbsTransport implements ConnectionMonitor, InteractiveC
 
 			// save this key in-memory if option enabled
 			if(manager.isSavingKeys()) {
-				manager.addKey(pubkey.getNickname(), trileadKey);
+				manager.addKey(pubkey, trileadKey);
 			}
 		}
 
@@ -861,18 +870,28 @@ public class SSH extends AbsTransport implements ConnectionMonitor, InteractiveC
 		if (useAuthAgent.equals(HostDatabase.AUTHAGENT_NO)) {
 			Log.e(TAG, "");
 			return null;
-		} else if (useAuthAgent.equals(HostDatabase.AUTHAGENT_CONFIRM)) {
-			Boolean result = bridge.promptHelper.requestBooleanPrompt(null,
-					manager.res.getString(R.string.prompt_allow_agent_to_use_key,
-							nickname));
-			if (result == null || !result)
+		} else if (useAuthAgent.equals(HostDatabase.AUTHAGENT_CONFIRM) ||
+				manager.loadedKeypairs.get(nickname).bean.isConfirmUse()) {
+			if (!promptForPubkeyUse(nickname))
 				return null;
 		}
 		return manager.getKey(nickname);
 	}
 
-	public boolean addIdentity(Object key, String comment) {
-		manager.addKey(comment, key);
+	private boolean promptForPubkeyUse(String nickname) {
+		Boolean result = bridge.promptHelper.requestBooleanPrompt(null,
+				manager.res.getString(R.string.prompt_allow_agent_to_use_key,
+						nickname));
+		return result;
+	}
+
+	public boolean addIdentity(Object key, String comment, boolean confirmUse, int lifetime) {
+		PubkeyBean pubkey = new PubkeyBean();
+//		pubkey.setType(PubkeyDatabase.KEY_TYPE_IMPORTED);
+		pubkey.setNickname(comment);
+		pubkey.setConfirmUse(confirmUse);
+		pubkey.setLifetime(lifetime);
+		manager.addKey(pubkey, key);
 		return true;
 	}
 
diff --git a/src/org/connectbot/util/PubkeyDatabase.java b/src/org/connectbot/util/PubkeyDatabase.java
index b7481b0..a813fa4 100644
--- a/src/org/connectbot/util/PubkeyDatabase.java
+++ b/src/org/connectbot/util/PubkeyDatabase.java
@@ -39,7 +39,7 @@ public class PubkeyDatabase extends RobustSQLiteOpenHelper {
 	public final static String TAG = "ConnectBot.PubkeyDatabase";
 
 	public final static String DB_NAME = "pubkeys";
-	public final static int DB_VERSION = 1;
+	public final static int DB_VERSION = 2;
 
 	public final static String TABLE_PUBKEYS = "pubkeys";
 	public final static String FIELD_PUBKEY_NICKNAME = "nickname";
@@ -48,6 +48,8 @@ public class PubkeyDatabase extends RobustSQLiteOpenHelper {
 	public final static String FIELD_PUBKEY_PUBLIC = "public";
 	public final static String FIELD_PUBKEY_ENCRYPTED = "encrypted";
 	public final static String FIELD_PUBKEY_STARTUP = "startup";
+	public final static String FIELD_PUBKEY_CONFIRMUSE = "confirmuse";
+	public final static String FIELD_PUBKEY_LIFETIME = "lifetime";
 
 	public final static String KEY_TYPE_RSA = "RSA",
 		KEY_TYPE_DSA = "DSA",
@@ -76,12 +78,20 @@ public class PubkeyDatabase extends RobustSQLiteOpenHelper {
 				+ FIELD_PUBKEY_PRIVATE + " BLOB, "
 				+ FIELD_PUBKEY_PUBLIC + " BLOB, "
 				+ FIELD_PUBKEY_ENCRYPTED + " INTEGER, "
-				+ FIELD_PUBKEY_STARTUP + " INTEGER)");
+				+ FIELD_PUBKEY_STARTUP + " INTEGER, "
+				+ FIELD_PUBKEY_CONFIRMUSE + " INTEGER DEFAULT 0, "
+				+ FIELD_PUBKEY_LIFETIME + " INTEGER DEFAULT 0)");
 	}
 
 	@Override
 	public void onRobustUpgrade(SQLiteDatabase db, int oldVersion, int newVersion) throws SQLiteException {
-
+			switch (oldVersion) {
+			case 1:
+				db.execSQL("ALTER TABLE " + TABLE_PUBKEYS
+						+ " ADD COLUMN " + FIELD_PUBKEY_CONFIRMUSE + " INTEGER DEFAULT 0");
+				db.execSQL("ALTER TABLE " + TABLE_PUBKEYS
+						+ " ADD COLUMN " + FIELD_PUBKEY_LIFETIME + " INTEGER DEFAULT 0");
+			}
 	}
 
 	/**
@@ -131,7 +141,9 @@ public class PubkeyDatabase extends RobustSQLiteOpenHelper {
 				COL_PRIVATE = c.getColumnIndexOrThrow(FIELD_PUBKEY_PRIVATE),
 				COL_PUBLIC = c.getColumnIndexOrThrow(FIELD_PUBKEY_PUBLIC),
 				COL_ENCRYPTED = c.getColumnIndexOrThrow(FIELD_PUBKEY_ENCRYPTED),
-				COL_STARTUP = c.getColumnIndexOrThrow(FIELD_PUBKEY_STARTUP);
+				COL_STARTUP = c.getColumnIndexOrThrow(FIELD_PUBKEY_STARTUP),
+				COL_CONFIRMUSE = c.getColumnIndexOrThrow(FIELD_PUBKEY_CONFIRMUSE),
+				COL_LIFETIME = c.getColumnIndexOrThrow(FIELD_PUBKEY_LIFETIME);
 
 			while (c.moveToNext()) {
 				PubkeyBean pubkey = new PubkeyBean();
@@ -143,6 +155,8 @@ public class PubkeyDatabase extends RobustSQLiteOpenHelper {
 				pubkey.setPublicKey(c.getBlob(COL_PUBLIC));
 				pubkey.setEncrypted(c.getInt(COL_ENCRYPTED) > 0);
 				pubkey.setStartup(c.getInt(COL_STARTUP) > 0);
+				pubkey.setConfirmUse(c.getInt(COL_CONFIRMUSE) > 0);
+				pubkey.setLifetime(c.getInt(COL_LIFETIME));
 
 				pubkeys.add(pubkey);
 			}
@@ -190,6 +204,8 @@ public class PubkeyDatabase extends RobustSQLiteOpenHelper {
 		pubkey.setPublicKey(c.getBlob(c.getColumnIndexOrThrow(FIELD_PUBKEY_PUBLIC)));
 		pubkey.setEncrypted(c.getInt(c.getColumnIndexOrThrow(FIELD_PUBKEY_ENCRYPTED)) > 0);
 		pubkey.setStartup(c.getInt(c.getColumnIndexOrThrow(FIELD_PUBKEY_STARTUP)) > 0);
+		pubkey.setConfirmUse(c.getInt(c.getColumnIndexOrThrow(FIELD_PUBKEY_CONFIRMUSE)) > 0);
+		pubkey.setLifetime(c.getInt(c.getColumnIndexOrThrow(FIELD_PUBKEY_LIFETIME)));
 
 		return pubkey;
 	}