1 files changed, 68 insertions, 9 deletions

diff --git a/example/ios/iOS UI Test/iOS UI Test/MasterViewController.m b/example/ios/iOS UI Test/iOS UI Test/MasterViewController.m
index 3e5e4422..5bbc0d53 100644
--- a/example/ios/iOS UI Test/iOS UI Test/MasterViewController.m
+++ b/example/ios/iOS UI Test/iOS UI Test/MasterViewController.m
@@ -10,6 +10,11 @@
 #import <MailCore/MailCore.h>
 #import "FXKeychain.h"
 #import "MCTMsgViewController.h"
+#import "GTMOAuth2ViewControllerTouch.h"
+
+#define CLIENT_ID @"the-client-id"
+#define CLIENT_SECRET @"the-client-secret"
+#define KEYCHAIN_ITEM_NAME @"MailCore OAuth 2.0 Token"
 
 @interface MasterViewController ()
 @property (nonatomic, strong) NSArray *messages;
@@ -26,24 +31,78 @@
 	
 	[[NSUserDefaults standardUserDefaults] registerDefaults:@{ HostnameKey: @"imap.gmail.com" }];
 	
+    if ([[NSUserDefaults standardUserDefaults] boolForKey:@"OAuth2Enabled"]) {
+        [self startOAuth2];
+    }
+    else {
+        [self startLogin];
+    }
+}
+
+- (void) startLogin
+{
 	NSString *username = [[NSUserDefaults standardUserDefaults] objectForKey:UsernameKey];
 	NSString *password = [[FXKeychain defaultKeychain] objectForKey:PasswordKey];
 	NSString *hostname = [[NSUserDefaults standardUserDefaults] objectForKey:HostnameKey];
-	
-	[self loadAccountWithUsername:username password:password hostname:hostname];
+    
+    if (!username.length || !password.length) {
+        [self performSelector:@selector(showSettingsViewController:) withObject:nil afterDelay:0.5];
+        return;
+    }
+
+	[self loadAccountWithUsername:username password:password hostname:hostname oauth2Token:nil];
 }
 
-- (void)loadAccountWithUsername:(NSString *)username password:(NSString *)password hostname:(NSString *)hostname {
-	if (!username.length || !password.length) {
-		[self performSelector:@selector(showSettingsViewController:) withObject:nil afterDelay:0.5];
-		return;
-	}
-	
+- (void) startOAuth2
+{
+    GTMOAuth2Authentication * auth = [GTMOAuth2ViewControllerTouch authForGoogleFromKeychainForName:KEYCHAIN_ITEM_NAME
+                                                                                        clientID:CLIENT_ID
+                                                                                    clientSecret:CLIENT_SECRET];
+    
+    if ([auth refreshToken] == nil) {
+        MasterViewController * __weak weakSelf = self;
+        GTMOAuth2ViewControllerTouch *viewController = [GTMOAuth2ViewControllerTouch controllerWithScope:@"https://mail.google.com/"
+                                                                                               clientID:CLIENT_ID
+                                                                                           clientSecret:CLIENT_SECRET
+                                                                                       keychainItemName:KEYCHAIN_ITEM_NAME
+                                                                                       completionHandler:^(GTMOAuth2ViewControllerTouch *viewController, GTMOAuth2Authentication *retrievedAuth, NSError *error) {
+                                                                                           [weakSelf loadWithAuth:retrievedAuth];
+                                                                                       }];
+        [self.navigationController pushViewController:viewController
+                                             animated:YES];
+    }
+    else {
+        [auth beginTokenFetchWithDelegate:self
+                        didFinishSelector:@selector(auth:finishedRefreshWithFetcher:error:)];
+    }
+}
+
+- (void)auth:(GTMOAuth2Authentication *)auth
+finishedRefreshWithFetcher:(GTMHTTPFetcher *)fetcher
+       error:(NSError *)error {
+    [self loadWithAuth:auth];
+}
+
+- (void)loadWithAuth:(GTMOAuth2Authentication *)auth
+{
+	NSString *hostname = [[NSUserDefaults standardUserDefaults] objectForKey:HostnameKey];
+	[self loadAccountWithUsername:[auth userEmail] password:nil hostname:hostname oauth2Token:[auth accessToken]];
+}
+
+- (void)loadAccountWithUsername:(NSString *)username
+                       password:(NSString *)password
+                       hostname:(NSString *)hostname
+                    oauth2Token:(NSString *)oauth2Token
+{
 	self.imapSession = [[MCOIMAPSession alloc] init];
 	self.imapSession.hostname = hostname;
 	self.imapSession.port = 993;
 	self.imapSession.username = username;
 	self.imapSession.password = password;
+    if (oauth2Token != nil) {
+        self.imapSession.OAuth2Token = oauth2Token;
+        self.imapSession.authType = MCOAuthTypeXOAuth2;
+    }
 	self.imapSession.connectionType = MCOConnectionTypeTLS;
     MasterViewController * __weak weakSelf = self;
 	self.imapSession.connectionLogger = ^(void * connectionID, MCOConnectionLogType type, NSData * data) {
@@ -135,7 +194,7 @@
 		![password isEqualToString:self.imapSession.password] ||
 		![hostname isEqualToString:self.imapSession.hostname]) {
 		self.imapSession = nil;
-		[self loadAccountWithUsername:username password:password hostname:hostname];
+		[self loadAccountWithUsername:username password:password hostname:hostname oauth2Token:nil];
 	}
 }