Two factor authentication support (#630)

* Initial commit for 2FA support

Signed-off-by: Andrew <write@imaginarycode.com>

* Add vendored files

* Add missing depends

* A few clean ups

* Added improvements, proper encryption

* Better encryption key

* Simplify "key" generation

* Make 2FA enrollment page more robust

* Fix typo

* Rename twofa/2FA to TwoFactor

* UNIQUE INDEX -> UNIQUE

2 files changed, 363 insertions, 8 deletions

diff --git a/routers/user/auth.go b/routers/user/auth.go
index ef2a04005b..85a51620ea 100644
--- a/routers/user/auth.go
+++ b/routers/user/auth.go
@@ -5,6 +5,7 @@
 package user
 
 import (
+	"errors"
 	"fmt"
 	"net/url"
 
@@ -27,6 +28,8 @@ const (
 	TplActivate       base.TplName = "user/auth/activate"
 	tplForgotPassword base.TplName = "user/auth/forgot_passwd"
 	tplResetPassword  base.TplName = "user/auth/reset_passwd"
+	tplTwofa          base.TplName = "user/auth/twofa"
+	tplTwofaScratch   base.TplName = "user/auth/twofa_scratch"
 )
 
 // AutoSignIn reads cookie and try to auto-login.
@@ -69,15 +72,12 @@ func AutoSignIn(ctx *context.Context) (bool, error) {
 	return true, nil
 }
 
-// SignIn render sign in page
-func SignIn(ctx *context.Context) {
-	ctx.Data["Title"] = ctx.Tr("sign_in")
-
+func checkAutoLogin(ctx *context.Context) bool {
 	// Check auto-login.
 	isSucceed, err := AutoSignIn(ctx)
 	if err != nil {
 		ctx.Handle(500, "AutoSignIn", err)
-		return
+		return true
 	}
 
 	redirectTo := ctx.Query("redirect_to")
@@ -94,6 +94,18 @@ func SignIn(ctx *context.Context) {
 		} else {
 			ctx.Redirect(setting.AppSubURL + "/")
 		}
+		return true
+	}
+
+	return false
+}
+
+// SignIn render sign in page
+func SignIn(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("sign_in")
+
+	// Check auto-login.
+	if checkAutoLogin(ctx) {
 		return
 	}
 
@@ -119,13 +131,158 @@ func SignInPost(ctx *context.Context, form auth.SignInForm) {
 		return
 	}
 
-	if form.Remember {
+	// If this user is enrolled in 2FA, we can't sign the user in just yet.
+	// Instead, redirect them to the 2FA authentication page.
+	_, err = models.GetTwoFactorByUID(u.ID)
+	if err != nil {
+		if models.IsErrTwoFactorNotEnrolled(err) {
+			handleSignIn(ctx, u, form.Remember)
+		} else {
+			ctx.Handle(500, "UserSignIn", err)
+		}
+		return
+	}
+
+	// User needs to use 2FA, save data and redirect to 2FA page.
+	ctx.Session.Set("twofaUid", u.ID)
+	ctx.Session.Set("twofaRemember", form.Remember)
+	ctx.Redirect(setting.AppSubURL + "/user/two_factor")
+}
+
+// TwoFactor shows the user a two-factor authentication page.
+func TwoFactor(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("twofa")
+
+	// Check auto-login.
+	if checkAutoLogin(ctx) {
+		return
+	}
+
+	// Ensure user is in a 2FA session.
+	if ctx.Session.Get("twofaUid") == nil {
+		ctx.Handle(500, "UserSignIn", errors.New("not in 2FA session"))
+		return
+	}
+
+	ctx.HTML(200, tplTwofa)
+}
+
+// TwoFactorPost validates a user's two-factor authentication token.
+func TwoFactorPost(ctx *context.Context, form auth.TwoFactorAuthForm) {
+	ctx.Data["Title"] = ctx.Tr("twofa")
+
+	// Ensure user is in a 2FA session.
+	idSess := ctx.Session.Get("twofaUid")
+	if idSess == nil {
+		ctx.Handle(500, "UserSignIn", errors.New("not in 2FA session"))
+		return
+	}
+
+	id := idSess.(int64)
+	twofa, err := models.GetTwoFactorByUID(id)
+	if err != nil {
+		ctx.Handle(500, "UserSignIn", err)
+		return
+	}
+
+	// Validate the passcode with the stored TOTP secret.
+	ok, err := twofa.ValidateTOTP(form.Passcode)
+	if err != nil {
+		ctx.Handle(500, "UserSignIn", err)
+		return
+	}
+
+	if ok {
+		remember := ctx.Session.Get("twofaRemember").(bool)
+		u, err := models.GetUserByID(id)
+		if err != nil {
+			ctx.Handle(500, "UserSignIn", err)
+			return
+		}
+
+		handleSignIn(ctx, u, remember)
+		return
+	}
+
+	ctx.RenderWithErr(ctx.Tr("auth.twofa_passcode_incorrect"), tplTwofa, auth.TwoFactorAuthForm{})
+}
+
+// TwoFactorScratch shows the scratch code form for two-factor authentication.
+func TwoFactorScratch(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("twofa_scratch")
+
+	// Check auto-login.
+	if checkAutoLogin(ctx) {
+		return
+	}
+
+	// Ensure user is in a 2FA session.
+	if ctx.Session.Get("twofaUid") == nil {
+		ctx.Handle(500, "UserSignIn", errors.New("not in 2FA session"))
+		return
+	}
+
+	ctx.HTML(200, tplTwofaScratch)
+}
+
+// TwoFactorScratchPost validates and invalidates a user's two-factor scratch token.
+func TwoFactorScratchPost(ctx *context.Context, form auth.TwoFactorScratchAuthForm) {
+	ctx.Data["Title"] = ctx.Tr("twofa_scratch")
+
+	// Ensure user is in a 2FA session.
+	idSess := ctx.Session.Get("twofaUid")
+	if idSess == nil {
+		ctx.Handle(500, "UserSignIn", errors.New("not in 2FA session"))
+		return
+	}
+
+	id := idSess.(int64)
+	twofa, err := models.GetTwoFactorByUID(id)
+	if err != nil {
+		ctx.Handle(500, "UserSignIn", err)
+		return
+	}
+
+	// Validate the passcode with the stored TOTP secret.
+	if twofa.VerifyScratchToken(form.Token) {
+		// Invalidate the scratch token.
+		twofa.ScratchToken = ""
+		if err = models.UpdateTwoFactor(twofa); err != nil {
+			ctx.Handle(500, "UserSignIn", err)
+			return
+		}
+
+		remember := ctx.Session.Get("twofaRemember").(bool)
+		u, err := models.GetUserByID(id)
+		if err != nil {
+			ctx.Handle(500, "UserSignIn", err)
+			return
+		}
+
+		handleSignInFull(ctx, u, remember, false)
+		ctx.Flash.Info(ctx.Tr("auth.twofa_scratch_used"))
+		ctx.Redirect(setting.AppSubURL + "/user/settings/two_factor")
+		return
+	}
+
+	ctx.RenderWithErr(ctx.Tr("auth.twofa_scratch_token_incorrect"), tplTwofaScratch, auth.TwoFactorScratchAuthForm{})
+}
+
+// This handles the final part of the sign-in process of the user.
+func handleSignIn(ctx *context.Context, u *models.User, remember bool) {
+	handleSignInFull(ctx, u, remember, true)
+}
+
+func handleSignInFull(ctx *context.Context, u *models.User, remember bool, obeyRedirect bool) {
+	if remember {
 		days := 86400 * setting.LogInRememberDays
 		ctx.SetCookie(setting.CookieUserName, u.Name, days, setting.AppSubURL)
 		ctx.SetSuperSecureCookie(base.EncodeMD5(u.Rands+u.Passwd),
 			setting.CookieRememberName, u.Name, days, setting.AppSubURL)
 	}
 
+	ctx.Session.Delete("twofaUid")
+	ctx.Session.Delete("twofaRemember")
 	ctx.Session.Set("uid", u.ID)
 	ctx.Session.Set("uname", u.Name)
 
@@ -141,11 +298,15 @@ func SignInPost(ctx *context.Context, form auth.SignInForm) {
 
 	if redirectTo, _ := url.QueryUnescape(ctx.GetCookie("redirect_to")); len(redirectTo) > 0 {
 		ctx.SetCookie("redirect_to", "", -1, setting.AppSubURL)
-		ctx.Redirect(redirectTo)
+		if obeyRedirect {
+			ctx.Redirect(redirectTo)
+		}
 		return
 	}
 
-	ctx.Redirect(setting.AppSubURL + "/")
+	if obeyRedirect {
+		ctx.Redirect(setting.AppSubURL + "/")
+	}
 }
 
 // SignOut sign out from login status
diff --git a/routers/user/setting.go b/routers/user/setting.go
index bba0f1e876..8fc2be0ce4 100644
--- a/routers/user/setting.go
+++ b/routers/user/setting.go
@@ -5,12 +5,19 @@
 package user
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"io/ioutil"
 	"strings"
 
 	"github.com/Unknwon/com"
+	"github.com/pquerna/otp"
+	"github.com/pquerna/otp/totp"
+
+	"encoding/base64"
+	"html/template"
+	"image/png"
 
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/auth"
@@ -28,6 +35,8 @@ const (
 	tplSettingsSSHKeys      base.TplName = "user/settings/sshkeys"
 	tplSettingsSocial       base.TplName = "user/settings/social"
 	tplSettingsApplications base.TplName = "user/settings/applications"
+	tplSettingsTwofa        base.TplName = "user/settings/twofa"
+	tplSettingsTwofaEnroll  base.TplName = "user/settings/twofa_enroll"
 	tplSettingsDelete       base.TplName = "user/settings/delete"
 	tplSecurity             base.TplName = "user/security"
 )
@@ -437,6 +446,191 @@ func SettingsDeleteApplication(ctx *context.Context) {
 	})
 }
 
+// SettingsTwoFactor renders the 2FA page.
+func SettingsTwoFactor(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("settings")
+	ctx.Data["PageIsSettingsTwofa"] = true
+
+	enrolled := true
+	_, err := models.GetTwoFactorByUID(ctx.User.ID)
+	if err != nil {
+		if models.IsErrTwoFactorNotEnrolled(err) {
+			enrolled = false
+		} else {
+			ctx.Handle(500, "SettingsTwoFactor", err)
+			return
+		}
+	}
+
+	ctx.Data["TwofaEnrolled"] = enrolled
+	ctx.HTML(200, tplSettingsTwofa)
+}
+
+// SettingsTwoFactorRegenerateScratch regenerates the user's 2FA scratch code.
+func SettingsTwoFactorRegenerateScratch(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("settings")
+	ctx.Data["PageIsSettingsTwofa"] = true
+
+	t, err := models.GetTwoFactorByUID(ctx.User.ID)
+	if err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	if err = t.GenerateScratchToken(); err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	if err = models.UpdateTwoFactor(t); err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	ctx.Flash.Success(ctx.Tr("settings.twofa_scratch_token_regenerated", t.ScratchToken))
+	ctx.Redirect(setting.AppSubURL + "/user/settings/two_factor")
+}
+
+// SettingsTwoFactorDisable deletes the user's 2FA settings.
+func SettingsTwoFactorDisable(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("settings")
+	ctx.Data["PageIsSettingsTwofa"] = true
+
+	t, err := models.GetTwoFactorByUID(ctx.User.ID)
+	if err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	if err = models.DeleteTwoFactorByID(t.ID, ctx.User.ID); err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	ctx.Flash.Success(ctx.Tr("settings.twofa_disabled"))
+	ctx.Redirect(setting.AppSubURL + "/user/settings/two_factor")
+}
+
+func twofaGenerateSecretAndQr(ctx *context.Context) bool {
+	var otpKey *otp.Key
+	var err error
+	uri := ctx.Session.Get("twofaUri")
+	if uri != nil {
+		otpKey, err = otp.NewKeyFromURL(uri.(string))
+	}
+	if otpKey == nil {
+		err = nil // clear the error, in case the URL was invalid
+		otpKey, err = totp.Generate(totp.GenerateOpts{
+			Issuer:      setting.AppName,
+			AccountName: ctx.User.Name,
+		})
+		if err != nil {
+			ctx.Handle(500, "SettingsTwoFactor", err)
+			return false
+		}
+	}
+
+	ctx.Data["TwofaSecret"] = otpKey.Secret()
+	img, err := otpKey.Image(320, 240)
+	if err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return false
+	}
+
+	var imgBytes bytes.Buffer
+	if err = png.Encode(&imgBytes, img); err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return false
+	}
+
+	ctx.Data["QrUri"] = template.URL("data:image/png;base64," + base64.StdEncoding.EncodeToString(imgBytes.Bytes()))
+	ctx.Session.Set("twofaSecret", otpKey.Secret())
+	ctx.Session.Set("twofaUri", otpKey.String())
+	return true
+}
+
+// SettingsTwoFactorEnroll shows the page where the user can enroll into 2FA.
+func SettingsTwoFactorEnroll(ctx *context.Context) {
+	ctx.Data["Title"] = ctx.Tr("settings")
+	ctx.Data["PageIsSettingsTwofa"] = true
+
+	t, err := models.GetTwoFactorByUID(ctx.User.ID)
+	if t != nil {
+		// already enrolled
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+	if err != nil && !models.IsErrTwoFactorNotEnrolled(err) {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	if !twofaGenerateSecretAndQr(ctx) {
+		return
+	}
+
+	ctx.HTML(200, tplSettingsTwofaEnroll)
+}
+
+// SettingsTwoFactorEnrollPost handles enrolling the user into 2FA.
+func SettingsTwoFactorEnrollPost(ctx *context.Context, form auth.TwoFactorAuthForm) {
+	ctx.Data["Title"] = ctx.Tr("settings")
+	ctx.Data["PageIsSettingsTwofa"] = true
+
+	t, err := models.GetTwoFactorByUID(ctx.User.ID)
+	if t != nil {
+		// already enrolled
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+	if err != nil && !models.IsErrTwoFactorNotEnrolled(err) {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	if ctx.HasError() {
+		if !twofaGenerateSecretAndQr(ctx) {
+			return
+		}
+		ctx.HTML(200, tplSettingsTwofaEnroll)
+		return
+	}
+
+	secret := ctx.Session.Get("twofaSecret").(string)
+	if !totp.Validate(form.Passcode, secret) {
+		if !twofaGenerateSecretAndQr(ctx) {
+			return
+		}
+		ctx.Flash.Error(ctx.Tr("settings.passcode_invalid"))
+		ctx.HTML(200, tplSettingsTwofaEnroll)
+		return
+	}
+
+	t = &models.TwoFactor{
+		UID: ctx.User.ID,
+	}
+	err = t.SetSecret(secret)
+	if err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+	err = t.GenerateScratchToken()
+	if err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	if err = models.NewTwoFactor(t); err != nil {
+		ctx.Handle(500, "SettingsTwoFactor", err)
+		return
+	}
+
+	ctx.Session.Delete("twofaSecret")
+	ctx.Session.Delete("twofaUri")
+	ctx.Flash.Success(ctx.Tr("settings.twofa_enrolled", t.ScratchToken))
+	ctx.Redirect(setting.AppSubURL + "/user/settings/two_factor")
+}
+
 // SettingsDelete render user suicide page and response for delete user himself
 func SettingsDelete(ctx *context.Context) {
 	ctx.Data["Title"] = ctx.Tr("settings")