From 15dcb43f44793950b5cde79d78353189a6baae75 Mon Sep 17 00:00:00 2001
From: Kyle Spearrin <kyle.spearrin@gmail.com>
Date: Fri, 23 Jun 2017 10:08:29 -0400
Subject: [PATCH] remember two factor token

---
 src/Core/Enums/TwoFactorProviderType.cs       |  3 ++-
 .../TwoFactorRememberTokenProvider.cs         | 21 +++++++++++++++++++
 .../ResourceOwnerPasswordValidator.cs         | 14 +++++++++++--
 .../Utilities/ServiceCollectionExtensions.cs  |  7 +++++++
 4 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 src/Core/Identity/TwoFactorRememberTokenProvider.cs

diff --git a/src/Core/Enums/TwoFactorProviderType.cs b/src/Core/Enums/TwoFactorProviderType.cs
index a23712a912..28be1b2c1b 100644
--- a/src/Core/Enums/TwoFactorProviderType.cs
+++ b/src/Core/Enums/TwoFactorProviderType.cs
@@ -6,6 +6,7 @@
         Email = 1,
         Duo = 2,
         YubiKey = 3,
-        U2f = 4
+        U2f = 4,
+        Remember = 5
     }
 }
diff --git a/src/Core/Identity/TwoFactorRememberTokenProvider.cs b/src/Core/Identity/TwoFactorRememberTokenProvider.cs
new file mode 100644
index 0000000000..f7ee1d07a9
--- /dev/null
+++ b/src/Core/Identity/TwoFactorRememberTokenProvider.cs
@@ -0,0 +1,21 @@
+using Microsoft.AspNetCore.Identity;
+using Bit.Core.Models.Table;
+using Microsoft.Extensions.Options;
+using Microsoft.AspNetCore.DataProtection;
+
+namespace Bit.Core.Identity
+{
+    public class TwoFactorRememberTokenProvider : DataProtectorTokenProvider<User>
+    {
+        private readonly GlobalSettings _globalSettings;
+
+        public TwoFactorRememberTokenProvider(
+            IDataProtectionProvider dataProtectionProvider,
+            IOptions<TwoFactorRememberTokenProviderOptions> options)
+            : base(dataProtectionProvider, options)
+        { }
+    }
+
+    public class TwoFactorRememberTokenProviderOptions : DataProtectionTokenProviderOptions
+    { }
+}
diff --git a/src/Core/IdentityServer/ResourceOwnerPasswordValidator.cs b/src/Core/IdentityServer/ResourceOwnerPasswordValidator.cs
index ba7bf0131a..0a1d111185 100644
--- a/src/Core/IdentityServer/ResourceOwnerPasswordValidator.cs
+++ b/src/Core/IdentityServer/ResourceOwnerPasswordValidator.cs
@@ -38,6 +38,7 @@ namespace Bit.Core.IdentityServer
         {
             var twoFactorToken = context.Request.Raw["TwoFactorToken"]?.ToString();
             var twoFactorProvider = context.Request.Raw["TwoFactorProvider"]?.ToString();
+            var twoFactorRemember = context.Request.Raw["TwoFactorRemember"]?.ToString() == "1";
             var twoFactorRequest = !string.IsNullOrWhiteSpace(twoFactorToken) && !string.IsNullOrWhiteSpace(twoFactorProvider);
 
             if(!string.IsNullOrWhiteSpace(context.UserName))
@@ -63,7 +64,8 @@ namespace Bit.Core.IdentityServer
                         if(!twoFactorRequest || await VerifyTwoFactor(user, twoFactorProviderType, twoFactorToken))
                         {
                             var device = await SaveDeviceAsync(user, context);
-                            BuildSuccessResult(user, context, device);
+                            await BuildSuccessResultAsync(user, context, device, twoFactorRequest,
+                                twoFactorProviderType, twoFactorRemember);
                             return;
                         }
                     }
@@ -74,7 +76,8 @@ namespace Bit.Core.IdentityServer
             BuildErrorResult(twoFactorRequest, context);
         }
 
-        private void BuildSuccessResult(User user, ResourceOwnerPasswordValidationContext context, Device device)
+        private async Task BuildSuccessResultAsync(User user, ResourceOwnerPasswordValidationContext context, Device device,
+            bool twoFactorRequest, TwoFactorProviderType twoFactorProviderType, bool twoFactorRemember)
         {
             var claims = new List<Claim>();
 
@@ -94,6 +97,12 @@ namespace Bit.Core.IdentityServer
                 customResponse.Add("Key", user.Key);
             }
 
+            if(twoFactorRequest && twoFactorRemember)
+            {
+                var token = await _userManager.GenerateTwoFactorTokenAsync(user, TwoFactorProviderType.Remember.ToString());
+                customResponse.Add("TwoFactorToken", token);
+            }
+
             context.Result = new GrantValidationResult(user.Id.ToString(), "Application",
                 identityProvider: "bitwarden",
                 claims: claims.Count > 0 ? claims : null,
@@ -167,6 +176,7 @@ namespace Bit.Core.IdentityServer
                 case TwoFactorProviderType.Duo:
                 case TwoFactorProviderType.YubiKey:
                 case TwoFactorProviderType.U2f:
+                case TwoFactorProviderType.Remember:
                     return await _userManager.VerifyTwoFactorTokenAsync(user, type.ToString(), token);
                 case TwoFactorProviderType.Email:
                     return await _userService.VerifyTwoFactorEmailAsync(user, token);
diff --git a/src/Core/Utilities/ServiceCollectionExtensions.cs b/src/Core/Utilities/ServiceCollectionExtensions.cs
index b5a749d0e2..63c31387b4 100644
--- a/src/Core/Utilities/ServiceCollectionExtensions.cs
+++ b/src/Core/Utilities/ServiceCollectionExtensions.cs
@@ -14,6 +14,7 @@ using Microsoft.AspNetCore.Identity;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.WindowsAzure.Storage;
+using System;
 using SqlServerRepos = Bit.Core.Repositories.SqlServer;
 
 namespace Bit.Core.Utilities
@@ -71,6 +72,11 @@ namespace Bit.Core.Utilities
         {
             services.AddTransient<ILookupNormalizer, LowerInvariantLookupNormalizer>();
 
+            services.Configure<TwoFactorRememberTokenProviderOptions>(options =>
+            {
+                options.TokenLifespan = TimeSpan.FromDays(30);
+            });
+
             var identityBuilder = services.AddIdentity<User, Role>(options =>
             {
                 options.User = new UserOptions
@@ -102,6 +108,7 @@ namespace Bit.Core.Utilities
                 .AddTokenProvider<YubicoOtpTokenProvider>(TwoFactorProviderType.YubiKey.ToString())
                 .AddTokenProvider<DuoWebTokenProvider>(TwoFactorProviderType.Duo.ToString())
                 .AddTokenProvider<U2fTokenProvider>(TwoFactorProviderType.U2f.ToString())
+                .AddTokenProvider<TwoFactorRememberTokenProvider>(TwoFactorProviderType.Remember.ToString())
                 .AddTokenProvider<EmailTokenProvider<User>>(TokenOptions.DefaultEmailProvider);
 
             return identityBuilder;