- Add support for SSH identity files (public keys) in custom locations

 - Update localisable strings
 - Fix initial window position

7 files changed, 140 insertions, 4 deletions

diff --git a/Source/SPConnectionController.h b/Source/SPConnectionController.h
index 4c8a0a9d..67bdb77b 100644
--- a/Source/SPConnectionController.h
+++ b/Source/SPConnectionController.h
@@ -80,8 +80,10 @@
 	NSString *sshHost;
 	NSString *sshUser;
 	NSString *sshPassword;
+	int sshKeyLocationEnabled;
+	NSString *sshKeyLocation;
 	NSString *sshPort;
-@private NSString *favoritesPBoardType;
+	@private NSString *favoritesPBoardType;
 
 	NSString *connectionKeychainItemName;
 	NSString *connectionKeychainItemAccount;
@@ -109,6 +111,7 @@
 	IBOutlet NSSecureTextField *socketPasswordField;
 	IBOutlet NSSecureTextField *sshPasswordField;
 	IBOutlet NSSecureTextField *sshSSHPasswordField;
+	IBOutlet NSButton *sshSSHKeyButton;
 
 	IBOutlet NSButton *addToFavoritesButton;
 	IBOutlet NSButton *connectButton;
@@ -135,6 +138,8 @@
 @property (readwrite, retain) NSString *sshHost;
 @property (readwrite, retain) NSString *sshUser;
 @property (readwrite, retain) NSString *sshPassword;
+@property (readwrite, assign) int sshKeyLocationEnabled;
+@property (readwrite, retain) NSString *sshKeyLocation;
 @property (readwrite, retain) NSString *sshPort;
 
 @property (readwrite, retain) NSString *connectionKeychainItemName;
@@ -158,6 +163,7 @@
 - (void)addConnectionToDocument;
 
 // Interface interaction
+- (IBAction)chooseSSHKey:(id)sender;
 - (IBAction)editFavorites:(id)sender;
 - (IBAction)showHelp:(id)sender;
 - (void)resizeTabViewToConnectionType:(NSUInteger)theType animating:(BOOL)animate;
diff --git a/Source/SPConnectionController.m b/Source/SPConnectionController.m
index 9c7abe57..7937ef99 100644
--- a/Source/SPConnectionController.m
+++ b/Source/SPConnectionController.m
@@ -54,6 +54,8 @@
 @synthesize sshHost;
 @synthesize sshUser;
 @synthesize sshPassword;
+@synthesize sshKeyLocationEnabled;
+@synthesize sshKeyLocation;
 @synthesize sshPort;
 
 @synthesize connectionKeychainItemName;
@@ -181,6 +183,15 @@
 		return;
 	}
 
+	// If an SSH key has been provided, verify it exists
+	if ([self type] == SPSSHTunnelConnection && sshKeyLocationEnabled && sshKeyLocation) {
+		if (![[NSFileManager defaultManager] fileExistsAtPath:[sshKeyLocation stringByExpandingTildeInPath]]) {
+			[self setSshKeyLocationEnabled:NSOffState];
+			SPBeginAlertSheet(NSLocalizedString(@"SSH Key not found", @"SSH key check error"), NSLocalizedString(@"OK", @"OK button"), nil, nil, [tableDocument parentWindow], self, nil, nil, NSLocalizedString(@"A SSH key location was specified, but no file was found in the specified location.  Please re-select the key and try again.", @"SSH key not found message"));
+			return;
+		}
+	}
+
 	// Ensure that a socket connection is not inadvertently used
 	if (![self checkHost]) return;
 	
@@ -283,6 +294,11 @@
 		[sshTunnel setPassword:[self sshPassword]];
 	}
 
+	// Set the public key path if appropriate
+	if (sshKeyLocationEnabled && sshKeyLocation) {
+		[sshTunnel setKeyFilePath:sshKeyLocation];
+	}
+
 	// Set the callback function on the tunnel
 	[sshTunnel setConnectionStateChangeSelector:@selector(sshTunnelCallback:) delegate:self];
 
@@ -461,6 +477,47 @@
 #pragma mark Interface interaction
 
 /**
+ * Opens the SSH key selection window, ready to select a SSH key.
+ */
+- (IBAction)chooseSSHKey:(id)sender
+{
+	[favoritesTable deselectAll:self];
+	NSString *directoryPath = nil;
+	NSString *filePath = nil;
+
+	// If the custom key location is currently disabled - after the button
+	// action - leave it disabled and return without showing the sheet.
+	if (!sshKeyLocationEnabled) {
+		return;
+	}
+
+	// Otherwise open a panel at the last or default location
+	if (sshKeyLocation && [sshKeyLocation length]) {
+		filePath = [sshKeyLocation lastPathComponent];
+		directoryPath = [sshKeyLocation stringByDeletingLastPathComponent];
+	}
+	[[NSOpenPanel openPanel] beginSheetForDirectory:directoryPath
+											   file:filePath
+											  types:[NSArray arrayWithObjects:@"pem", @"", nil]
+									 modalForWindow:[tableDocument parentWindow]
+									  modalDelegate:self
+									 didEndSelector:@selector(chooseSSHKeySheetDidEnd:returnCode:contextInfo:)
+										contextInfo:nil];
+}
+
+/**
+ * Called after closing the SSH key selection sheet.
+ */
+- (void)chooseSSHKeySheetDidEnd:(NSOpenPanel *)openPanel returnCode:(NSInteger)returnCode contextInfo:(void *)contextInfo
+{
+	if (returnCode == NSCancelButton) {
+		[self setSshKeyLocationEnabled:NSOffState];
+		return;
+	}
+	[self setSshKeyLocation:[[openPanel filename] stringByAbbreviatingWithTildeInPath]];
+}
+
+/**
  * Opens the preferences window, or brings it to the front, and switch to the favorites tab.
  * If a favorite is selected in the connection sheet, it is also select in the prefs window.
  */
@@ -679,6 +736,8 @@
 	[self setDatabase:([self valueForKeyPath:@"selectedFavorite.database"] ? [self valueForKeyPath:@"selectedFavorite.database"] : @"")];
 	[self setSshHost:([self valueForKeyPath:@"selectedFavorite.sshHost"] ? [self valueForKeyPath:@"selectedFavorite.sshHost"] : @"")];
 	[self setSshUser:([self valueForKeyPath:@"selectedFavorite.sshUser"] ? [self valueForKeyPath:@"selectedFavorite.sshUser"] : @"")];
+	[self setSshKeyLocationEnabled:([self valueForKeyPath:@"selectedFavorite.sshKeyLocationEnabled"] ? [[self valueForKeyPath:@"selectedFavorite.sshKeyLocationEnabled"] intValue] : NSOffState)];
+	[self setSshKeyLocation:([self valueForKeyPath:@"selectedFavorite.sshKeyLocation"] ? [self valueForKeyPath:@"selectedFavorite.sshKeyLocation"] : @"")];
 	[self setSshPort:([self valueForKeyPath:@"selectedFavorite.sshPort"] ? [self valueForKeyPath:@"selectedFavorite.sshPort"] : @"")];
 
 	// Check whether the password exists in the keychain, and if so add it; also record the
@@ -773,6 +832,8 @@
 	if ([self database]) [newFavorite setObject:[self database] forKey:@"database"];
 	if ([self sshHost]) [newFavorite setObject:[self sshHost] forKey:@"sshHost"];
 	if ([self sshUser]) [newFavorite setObject:[self sshUser] forKey:@"sshUser"];
+	[newFavorite setObject:[NSNumber numberWithInt:[self sshKeyLocationEnabled]] forKey:@"sshKeyLocationEnabled"];
+	if ([self sshKeyLocation]) [newFavorite setObject:[self sshKeyLocation] forKey:@"sshKeyLocation"];
 	if ([self sshPort]) [newFavorite setObject:[self sshPort] forKey:@"sshPort"];
 
 	// Add the new favorite to the user defaults array
diff --git a/Source/SPDatabaseDocument.m b/Source/SPDatabaseDocument.m
index 7632d50a..47f915c3 100644
--- a/Source/SPDatabaseDocument.m
+++ b/Source/SPDatabaseDocument.m
@@ -268,6 +268,8 @@
 	[connectionController setSocket:@""];
 	[connectionController setSshHost:@""];
 	[connectionController setSshUser:@""];
+	[connectionController setSshKeyLocationEnabled:NSOffState];
+	[connectionController setSshKeyLocation:@""];
 	[connectionController setSshPort:@""];
 	[connectionController setDatabase:@""];
 	[connectionController setPassword:nil];
@@ -485,6 +487,10 @@
 			[connectionController setSshHost:[connection objectForKey:@"ssh_host"]];
 		if([connection objectForKey:@"ssh_user"])
 			[connectionController setSshUser:[connection objectForKey:@"ssh_user"]];
+		if([connection objectForKey:@"ssh_keyLocationEnabled"])
+			[connectionController setSshKeyLocationEnabled:[[connection objectForKey:@"ssh_keyLocationEnabled"] intValue]];
+		if([connection objectForKey:@"ssh_keyLocation"])
+			[connectionController setSshKeyLocation:[connection objectForKey:@"ssh_keyLocation"]];
 		if([connection objectForKey:@"ssh_port"])
 			[connectionController setSshPort:[NSString stringWithFormat:@"%ld", (long)[[connection objectForKey:@"ssh_port"] integerValue]]];
 
@@ -3293,6 +3299,8 @@
 		aString = @"SPSSHTunnelConnection";
 		[connection setObject:[connectionController sshHost] forKey:@"ssh_host"];
 		[connection setObject:[connectionController sshUser] forKey:@"ssh_user"];
+		[connection setObject:[NSNumber numberWithInt:[connectionController sshKeyLocationEnabled]] forKey:@"ssh_keyLocationEnabled"];
+		[connection setObject:[connectionController sshKeyLocation] forKey:@"ssh_keyLocation"];
 		if([connectionController sshPort] && [[connectionController sshPort] length])
 			[connection setObject:[NSNumber numberWithInteger:[[connectionController sshPort] integerValue]] forKey:@"ssh_port"];
 		break;
diff --git a/Source/SPPreferenceController.h b/Source/SPPreferenceController.h
index 39c5f58e..1683db65 100644
--- a/Source/SPPreferenceController.h
+++ b/Source/SPPreferenceController.h
@@ -132,7 +132,7 @@
 - (IBAction)closePanelSheet:(id)sender;
 - (IBAction)duplicateTheme:(id)sender;
 - (IBAction)removeTheme:(id)sender;
-
+- (IBAction)chooseSSHKey:(id)sender;
 
 // Toolbar item IBAction methods
 - (IBAction)displayGeneralPreferences:(id)sender;
diff --git a/Source/SPPreferenceController.m b/Source/SPPreferenceController.m
index 9f09ac2c..42503649 100644
--- a/Source/SPPreferenceController.m
+++ b/Source/SPPreferenceController.m
@@ -409,8 +409,8 @@
 	NSNumber *favoriteid = [NSNumber numberWithInteger:[[NSString stringWithFormat:@"%f", [[NSDate date] timeIntervalSince1970]] hash]];
 
 	// Create default favorite
-	NSMutableDictionary *favorite = [NSMutableDictionary dictionaryWithObjects:[NSArray arrayWithObjects:@"New Favorite", [NSNumber numberWithInteger:0], @"", @"", @"", @"", @"", @"", @"", @"", favoriteid, nil] 
-																	   forKeys:[NSArray arrayWithObjects:@"name", @"type", @"host", @"socket", @"user", @"port", @"database", @"sshHost", @"sshUser", @"sshPort", @"id", nil]];
+	NSMutableDictionary *favorite = [NSMutableDictionary dictionaryWithObjects:[NSArray arrayWithObjects:@"New Favorite", [NSNumber numberWithInteger:0], @"", @"", @"", @"", @"", @"", @"", [NSNumber numberWithInt:NSOffState], @"", @"", favoriteid, nil] 
+																	   forKeys:[NSArray arrayWithObjects:@"name", @"type", @"host", @"socket", @"user", @"port", @"database", @"sshHost", @"sshUser", @"sshKeyLocationEnabled", @"sshKeyLocation", @"sshPort", @"id", nil]];
 	
 	[favoritesController addObject:favorite];
 	[favoritesController setSelectedObjects:[NSArray arrayWithObject:favorite]];
@@ -685,6 +685,46 @@
 	[[sender window] orderOut:self];
 }
 
+/**
+ * Opens the SSH key selection window, ready to select a SSH key.
+ */
+- (IBAction)chooseSSHKey:(id)sender
+{
+	NSString *directoryPath = nil;
+	NSString *filePath = nil;
+
+	// If the custom key location is currently disabled - after the button
+	// action - leave it disabled and return without showing the sheet.
+	if (![[favoritesController valueForKeyPath:@"selection.sshKeyLocationEnabled"] intValue]) {
+		return;
+	}
+
+	// Otherwise open a panel at the last or default location
+	if ([favoritesController valueForKeyPath:@"selection.sshKeyLocation"] && [[favoritesController valueForKeyPath:@"selection.sshKeyLocation"] length]) {
+		filePath = [[favoritesController valueForKeyPath:@"selection.sshKeyLocation"] lastPathComponent];
+		directoryPath = [[favoritesController valueForKeyPath:@"selection.sshKeyLocation"] stringByDeletingLastPathComponent];
+	}
+	[[NSOpenPanel openPanel] beginSheetForDirectory:directoryPath
+											   file:filePath
+											  types:[NSArray arrayWithObjects:@"pem", @"", nil]
+									 modalForWindow:preferencesWindow
+									  modalDelegate:self
+									 didEndSelector:@selector(chooseSSHKeySheetDidEnd:returnCode:contextInfo:)
+										contextInfo:nil];
+}
+
+/**
+ * Called after closing the SSH key selection sheet.
+ */
+- (void)chooseSSHKeySheetDidEnd:(NSOpenPanel *)openPanel returnCode:(NSInteger)returnCode contextInfo:(void *)contextInfo
+{
+	if (returnCode == NSCancelButton) {
+		[favoritesController setValue:[NSNumber numberWithInt:NSOffState] forKeyPath:@"selection.sshKeyLocationEnabled"];
+		return;
+	}
+	[favoritesController setValue:[[openPanel filename] stringByAbbreviatingWithTildeInPath] forKeyPath:@"selection.sshKeyLocation"];
+}
+
 #pragma mark -
 #pragma mark Toolbar item IBAction methods
 
diff --git a/Source/SPSSHTunnel.h b/Source/SPSSHTunnel.h
index 9230f7e1..3af537c6 100644
--- a/Source/SPSSHTunnel.h
+++ b/Source/SPSSHTunnel.h
@@ -51,6 +51,7 @@
 	NSString *keychainName;
 	NSString *keychainAccount;
 	NSString *requestedPassphrase;
+	NSString *identityFilePath;
 	NSMutableArray *debugMessages;
 	NSLock *debugMessagesLock;
 	BOOL useHostFallback;
@@ -71,6 +72,7 @@
 - (void) setParentWindow:(NSWindow *)theWindow;
 - (BOOL) setPasswordKeychainName:(NSString *)theName account:(NSString *)theAccount;
 - (BOOL) setPassword:(NSString *)thePassword;
+- (BOOL) setKeyFilePath:(NSString *)thePath;
 - (NSInteger) state;
 - (NSString *) lastError;
 - (NSString *) debugMessages;
diff --git a/Source/SPSSHTunnel.m b/Source/SPSSHTunnel.m
index 5cc833fd..807d507e 100644
--- a/Source/SPSSHTunnel.m
+++ b/Source/SPSSHTunnel.m
@@ -72,6 +72,7 @@
 	}
 
 	parentWindow = nil;
+	identityFilePath = nil;
 	sshQuestionDialog = nil;
 	sshPasswordDialog = nil;
 	password = nil;
@@ -128,6 +129,19 @@
 	return YES;
 }
 
+/**
+ * Sets the path of an identity file, or public key file, to use when connecting.
+ */
+- (BOOL) setKeyFilePath:(NSString *)thePath
+{
+	NSString *expandedPath = [thePath stringByExpandingTildeInPath];
+	if (![[NSFileManager defaultManager] fileExistsAtPath:expandedPath]) return NO;
+
+	if (identityFilePath) [identityFilePath release];
+	identityFilePath = [[NSString alloc] initWithString:expandedPath];
+	return YES;
+}
+
 /*
  * Sets the keychain name to use to retrieve the password.  This is the recommended and
  * secure way of supplying a password to the SSH tunnel.
@@ -282,6 +296,10 @@
 	[taskArguments addObject:@"-o ExitOnForwardFailure=yes"];
 	[taskArguments addObject:[NSString stringWithFormat:@"-o ConnectTimeout=%ld", (long)connectionTimeout]];
 	[taskArguments addObject:@"-o NumberOfPasswordPrompts=3"];
+	if (identityFilePath) {
+		[taskArguments addObject:@"-i"];
+		[taskArguments addObject:identityFilePath];
+	}
 	if (useKeepAlive && keepAliveInterval) {
 		[taskArguments addObject:@"-o TCPKeepAlive=no"];		
 		[taskArguments addObject:[NSString stringWithFormat:@"-o ServerAliveInterval=%ld", (long)ceil(keepAliveInterval)]];		
@@ -652,6 +670,7 @@
 	if (password) [password release];
 	if (keychainName) [keychainName release];
 	if (keychainAccount) [keychainAccount release];
+	if (identityFilePath) [identityFilePath release];
 
 	// As this object is not a NSWindowController, use manual top-level nib item management
 	if (sshQuestionDialog) [sshQuestionDialog release], sshQuestionDialog = nil;