From 6ff2b3b729ee34bd5ca35bbd77360c3ac42796a3 Mon Sep 17 00:00:00 2001
From: Jingwen Owen Ou <jingweno@gmail.com>
Date: Tue, 15 Oct 2013 09:36:42 -0700
Subject: [PATCH] Implement two factor authentication

---
 github/config.go | 20 +++++++++++++++++++-
 github/github.go | 14 ++++++++++----
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/github/config.go b/github/config.go
index 5d432c79..d3397cc6 100644
--- a/github/config.go
+++ b/github/config.go
@@ -9,6 +9,7 @@ import (
 	"github.com/jingweno/gh/utils"
 	"os"
 	"path/filepath"
+	"regexp"
 )
 
 type Config struct {
@@ -40,6 +41,14 @@ func (c *Config) FetchPassword() string {
 	return string(pass)
 }
 
+func (c *Config) FetchTwoFactorCode() string {
+	var code string
+	fmt.Print("two-factor authentication code: ")
+	fmt.Scanln(&code)
+
+	return code
+}
+
 func (c *Config) FetchCredentials() {
 	var changed bool
 	if c.User == "" {
@@ -49,7 +58,16 @@ func (c *Config) FetchCredentials() {
 
 	if c.Token == "" {
 		password := c.FetchPassword()
-		token, err := findOrCreateToken(c.User, password)
+		token, err := findOrCreateToken(c.User, password, "")
+		// TODO: return an two factor auth failure error
+		if err != nil {
+			re := regexp.MustCompile("two-factor authentication OTP code")
+			if re.MatchString(fmt.Sprintf("%s", err)) {
+				code := c.FetchTwoFactorCode()
+				token, err = findOrCreateToken(c.User, password, code)
+			}
+		}
+
 		utils.Check(err)
 
 		c.Token = token
diff --git a/github/github.go b/github/github.go
index 5b1f76dd..91f3f8c4 100644
--- a/github/github.go
+++ b/github/github.go
@@ -127,9 +127,15 @@ func (gh *GitHub) repo() octokat.Repo {
 	return octokat.Repo{Name: project.Name, UserName: project.Owner}
 }
 
-func findOrCreateToken(user, password string) (string, error) {
+func findOrCreateToken(user, password, twoFactorCode string) (string, error) {
 	client := octokat.NewClient().WithLogin(user, password)
-	auths, err := client.Authorizations(nil)
+	options := &octokat.Options{}
+	if twoFactorCode != "" {
+		headers := octokat.Headers{"X-GitHub-OTP": twoFactorCode}
+		options.Headers = headers
+	}
+
+	auths, err := client.Authorizations(options)
 	if err != nil {
 		return "", err
 	}
@@ -147,9 +153,9 @@ func findOrCreateToken(user, password string) (string, error) {
 		authParam.Scopes = append(authParam.Scopes, "repo")
 		authParam.Note = "gh"
 		authParam.NoteURL = OAuthAppURL
-		options := octokat.Options{Params: authParam}
+		options.Params = authParam
 
-		auth, err := client.CreateAuthorization(&options)
+		auth, err := client.CreateAuthorization(options)
 		if err != nil {
 			return "", err
 		}
-- 
GitLab