✨ feat: totp close #9

.dockerignore

@@ -16,4 +16,5 @@
 # Dependency directories (remove the comment below to include it)
 # vendor/
 
+.gitignore
 web/node_modules/

.gitignore

@@ -0,0 +1,6 @@
+yarn.lock
+web/build
+*.log
+*.db
+.DS_Store
+.eslintcache

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/labstack/gommon v0.3.0
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/pkg/sftp v1.12.0
+	github.com/pquerna/otp v1.3.0
 	github.com/sirupsen/logrus v1.4.2
 	github.com/spf13/pflag v1.0.3
 	github.com/spf13/viper v1.7.1

go.sum

@@ -26,6 +26,8 @@ github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
 github.com/bketelsen/crypt v0.0.3-0.20200106085610-5cbc8cc4026c/go.mod h1:MKsuJmJgSg28kpZDP6UIiPt0e0Oz0kqKNGyRaWEPv84=
+github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc h1:biVzkmvwrH8WK8raXaxBx6fRVTlJILwEwQGL1I/ByEI=
+github.com/boombuler/barcode v1.0.1-0.20190219062509-6c824513bacc/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=
@@ -169,6 +171,8 @@ github.com/pkg/sftp v1.12.0/go.mod h1:fUqqXB5vEgVCZ131L+9say31RAri6aF6KDViawhxKK
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
+github.com/pquerna/otp v1.3.0 h1:oJV/SkzR33anKXwQU3Of42rL4wbrffP4uvUf1SvS5Xs=
+github.com/pquerna/otp v1.3.0/go.mod h1:dkJfzwRKNiegxyNb54X/3fLwhCynbMspSyWKnvi1AEg=
 github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
 github.com/prometheus/client_golang v0.9.3/go.mod h1:/TN21ttK/J9q6uSwhBd54HahCDft0ttaMvbicHlPoso=
 github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=

pkg/api/account.go

@@ -1,16 +1,26 @@
 package api
 
 import (
-	"github.com/labstack/echo/v4"
+	"log"
+	"time"
+
 	"next-terminal/pkg/global"
 	"next-terminal/pkg/model"
+	"next-terminal/pkg/totp"
 	"next-terminal/pkg/utils"
-	"time"
+
+	"github.com/labstack/echo/v4"
 )
 
 type LoginAccount struct {
 	Username string `json:"username"`
 	Password string `json:"password"`
+	TOTP     string `json:"totp"`
+}
+
+type ConfirmTOTP struct {
+	Secret string `json:"secret"`
+	TOTP   string `json:"totp"`
 }
 
 type ChangePassword struct {
@@ -28,10 +38,17 @@ func LoginEndpoint(c echo.Context) error {
 	if err != nil {
 		return Fail(c, -1, "您输入的账号或密码不正确")
 	}
+
 	if err := utils.Encoder.Match([]byte(user.Password), []byte(loginAccount.Password)); err != nil {
 		return Fail(c, -1, "您输入的账号或密码不正确")
 	}
 
+	log.Println(user, loginAccount)
+
+	if !totp.Validate(loginAccount.TOTP, user.TOTPSecret) {
+		return Fail(c, -2, "您的TOTP不匹配")
+	}
+
 	token := utils.UUID()
 
 	global.Cache.Set(token, user, time.Minute*time.Duration(30))
@@ -47,6 +64,54 @@ func LogoutEndpoint(c echo.Context) error {
 	return Success(c, nil)
 }
 
+func ConfirmTOTPEndpoint(c echo.Context) error {
+	account, _ := GetCurrentAccount(c)
+
+	var confirmTOTP ConfirmTOTP
+	if err := c.Bind(&confirmTOTP); err != nil {
+		return err
+	}
+
+	if !totp.Validate(confirmTOTP.TOTP, confirmTOTP.Secret) {
+		return Fail(c, -1, "TOTP 验证失败，请重试")
+	}
+
+	u := &model.User{
+		TOTPSecret: confirmTOTP.Secret,
+	}
+
+	model.UpdateUserById(u, account.ID)
+
+	return Success(c, nil)
+}
+
+func ResetTOTPEndpoint(c echo.Context) error {
+	account, _ := GetCurrentAccount(c)
+
+	key, err := totp.NewTOTP(totp.GenerateOpts{
+		Issuer:      c.Request().Host,
+		AccountName: account.Username,
+	})
+	if err != nil {
+		return Fail(c, -1, err.Error())
+	}
+
+	qrcode, err := key.Image(200, 200)
+	if err != nil {
+		return Fail(c, -1, err.Error())
+	}
+
+	qrEncode, err := utils.ImageToBase64Encode(qrcode)
+	if err != nil {
+		return Fail(c, -1, err.Error())
+	}
+
+	return Success(c, map[string]string{
+		"qr":     qrEncode,
+		"secret": key.Secret(),
+	})
+}
+
 func ChangePasswordEndpoint(c echo.Context) error {
 	account, _ := GetCurrentAccount(c)
 

pkg/api/routes.go

@@ -1,11 +1,12 @@
 package api
 
 import (
-	"github.com/labstack/echo/v4"
-	"github.com/labstack/echo/v4/middleware"
 	"net/http"
 	"next-terminal/pkg/global"
 	"next-terminal/pkg/model"
+
+	"github.com/labstack/echo/v4"
+	"github.com/labstack/echo/v4/middleware"
 )
 
 const Token = "X-Auth-Token"
@@ -35,6 +36,8 @@ func SetupRoutes() *echo.Echo {
 
 	e.POST("/logout", LogoutEndpoint)
 	e.POST("/change-password", ChangePasswordEndpoint)
+	e.POST("/reset-totp", ResetTOTPEndpoint)
+	e.POST("/confirm-totp", ConfirmTOTPEndpoint)
 	e.GET("/info", InfoEndpoint)
 
 	users := e.Group("/users")

pkg/model/user.go

@@ -11,6 +11,7 @@ type User struct {
 	Username   string         `json:"username"`
 	Password   string         `json:"password"`
 	Nickname   string         `json:"nickname"`
+	TOTPSecret string         `json:"-"`
 	Online     bool           `json:"online"`
 	Enabled    bool           `json:"enabled"`
 	Created    utils.JsonTime `json:"created"`
@@ -32,7 +33,6 @@ func FindAllUser() (o []User) {
 }
 
 func FindPageUser(pageIndex, pageSize int, username, nickname string) (o []User, total int64, err error) {
-
 	db := global.DB
 	if len(username) > 0 {
 		db = db.Where("username like ?", "%"+username+"%")

pkg/totp/totp.go

@@ -0,0 +1,19 @@
+package totp
+
+import (
+	otp_t "github.com/pquerna/otp"
+	totp_t "github.com/pquerna/otp/totp"
+)
+
+type GenerateOpts totp_t.GenerateOpts
+
+func NewTOTP(opt GenerateOpts) (*otp_t.Key, error) {
+	return totp_t.Generate(totp_t.GenerateOpts(opt))
+}
+
+func Validate(code string, secret string) bool {
+	if secret == "" {
+		return true
+	}
+	return totp_t.Validate(code, secret)
+}

pkg/utils/utils.go

@@ -1,14 +1,19 @@
 package utils
 
 import (
+	"bytes"
 	"database/sql/driver"
+	"encoding/base64"
 	"fmt"
-	"github.com/gofrs/uuid"
+	"image"
-	"golang.org/x/crypto/bcrypt"
+	"image/png"
 	"net"
 	"os"
 	"strconv"
 	"time"
+
+	"github.com/gofrs/uuid"
+	"golang.org/x/crypto/bcrypt"
 )
 
 type JsonTime struct {
@@ -81,6 +86,14 @@ func Tcping(ip string, port int) bool {
 	return true
 }
 
+func ImageToBase64Encode(img image.Image) (string, error) {
+	var buf bytes.Buffer
+	if err := png.Encode(&buf, img); err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(buf.Bytes()), nil
+}
+
 // 判断所给路径文件/文件夹是否存在
 func FileExists(path string) bool {
 	_, err := os.Stat(path) //os.Stat获取文件信息

web/src/components/Login.js

@@ -1,14 +1,12 @@
-import React, {Component} from 'react';
+import React, { Component } from 'react';
-import {Button, Card, Checkbox, Form, Input, Typography} from "antd";
+import { Button, Card, Checkbox, Form, Input, Typography } from "antd";
 import './Login.css'
 import request from "../common/request";
-import {message} from "antd/es";
+import { message } from "antd/es";
-import {withRouter} from "react-router-dom";
+import { withRouter } from "react-router-dom";
-import {LockOutlined, UserOutlined} from '@ant-design/icons';
+import { OneToOneOutlined, LockOutlined, UserOutlined } from '@ant-design/icons';
-
-
-const {Title} = Typography;
 
+const { Title } = Typography;
 class LoginForm extends Component {
 
     state = {
@@ -62,22 +60,23 @@ class LoginForm extends Component {
     render() {
         return (
             <div className='login-bg'
-                 style={{width: this.state.width, height: this.state.height, backgroundColor: '#F0F2F5'}}>
+                style={{ width: this.state.width, height: this.state.height, backgroundColor: '#F0F2F5' }}>
                 <Card className='login-card' title={null}>
-                    <div style={{textAlign: "center", margin: '15px auto 30px auto', color: '#1890ff'}}>
+                    <div style={{ textAlign: "center", margin: '15px auto 30px auto', color: '#1890ff' }}>
                         <Title level={1}>Next Terminal</Title>
                     </div>
                     <Form onFinish={this.handleSubmit} className="login-form">
-                        <Form.Item name='username' rules={[{required: true, message: '请输入登录账号！'}]}>
+                        <Form.Item name='username' rules={[{ required: true, message: '请输入登录账号！' }]}>
-                            <Input prefix={<UserOutlined/>} placeholder="登录账号"/>
+                            <Input prefix={<UserOutlined />} placeholder="登录账号" />
                         </Form.Item>
-                        <Form.Item name='password' rules={[{required: true, message: '请输入登录密码！'}]}>
+                        <Form.Item name='password' rules={[{ required: true, message: '请输入登录密码！' }]}>
-                            <Input.Password prefix={<LockOutlined/>} placeholder="登录密码"/>
+                            <Input.Password prefix={<LockOutlined />} placeholder="登录密码" />
+                        </Form.Item>
+                        <Form.Item name='totp' rules={[]}>
+                            <Input prefix={<OneToOneOutlined />} placeholder="TOTP" />
                         </Form.Item>
                         <Form.Item name='remember' valuePropName='checked' initialValue={false}>
                             <Checkbox>记住登录</Checkbox>
-
-
                         </Form.Item>
                         <Form.Item>
                             <Button type="primary" htmlType="submit" className="login-form-button"

web/src/components/user/Info.js

@@ -1,11 +1,11 @@
-import React, {Component} from 'react';
+import React, { Component } from 'react';
-import {Button, Form, Input, Layout, PageHeader} from "antd";
+import { Button, Form, Input, Layout, PageHeader, Image } from "antd";
-import {itemRender} from '../../utils/utils'
+import { itemRender } from '../../utils/utils'
 import request from "../../common/request";
-import {message} from "antd/es";
+import { message } from "antd/es";
 import Logout from "./Logout";
 
-const {Content} = Layout;
+const { Content } = Layout;
 
 const routes = [
     {
@@ -19,12 +19,12 @@ const routes = [
 ];
 
 const formItemLayout = {
-    labelCol: {span: 3},
+    labelCol: { span: 3 },
-    wrapperCol: {span: 6},
+    wrapperCol: { span: 6 },
 };
 const formTailLayout = {
-    labelCol: {span: 3},
+    labelCol: { span: 3 },
-    wrapperCol: {span: 6, offset: 3},
+    wrapperCol: { span: 6, offset: 3 },
 };
 
 class Info extends Component {
@@ -69,6 +69,32 @@ class Info extends Component {
         }
     }
 
+    confirmTOTP = async (values) => {
+        values['secret'] = this.state.secret
+        let result = await request.post('/confirm-totp', values);
+        if (result.code === 1) {
+            message.success('TOTP启用成功');
+            this.setState({
+                qr: "",
+                secret: ""
+            })
+        } else {
+            message.error(result.message);
+        }
+    }
+
+    resetTOTP = async () => {
+        let result = await request.post('/reset-totp');
+        if (result.code === 1) {
+            this.setState({
+                qr: result.data.qr,
+                secret: result.data.secret,
+            })
+        } else {
+            message.error(result.message);
+        }
+    }
+
     render() {
         return (
             <>
@@ -80,14 +106,12 @@ class Info extends Component {
                         itemRender: itemRender
                     }}
                     extra={[
-                        <Logout key='logout'/>
+                        <Logout key='logout' />
                     ]}
                     subTitle="个人中心"
-                >
+                />
-                </PageHeader>
 
                 <Content className="site-layout-background page-content">
-
                     <h1>修改密码</h1>
                     <Form ref={this.passwordFormRef} name="password" onFinish={this.changePassword}>
                         <Form.Item
@@ -101,7 +125,7 @@ class Info extends Component {
                                 },
                             ]}
                         >
-                            <Input type='password' placeholder="请输入原始密码"/>
+                            <Input type='password' placeholder="请输入原始密码" />
                         </Form.Item>
                         <Form.Item
                             {...formItemLayout}
@@ -115,9 +139,8 @@ class Info extends Component {
                             ]}
                         >
                             <Input type='password' placeholder="新的密码"
-                                   onChange={(value) => this.onNewPasswordChange(value)}/>
+                                onChange={(value) => this.onNewPasswordChange(value)} />
                         </Form.Item>
-
                         <Form.Item
                             {...formItemLayout}
                             name="newPassword2"
@@ -132,9 +155,8 @@ class Info extends Component {
                             help={this.state.errorMsg || ''}
                         >
                             <Input type='password' placeholder="请和上面输入新的密码保持一致"
-                                   onChange={(value) => this.onNewPassword2Change(value)}/>
+                                onChange={(value) => this.onNewPassword2Change(value)} />
                         </Form.Item>
-
                         <Form.Item {...formTailLayout}>
                             <Button type="primary" htmlType="submit">
                                 提交
@@ -142,6 +164,37 @@ class Info extends Component {
                         </Form.Item>
                     </Form>
                 </Content>
+                <Content className="site-layout-background page-content">
+                    <h1>双因素认证</h1>
+                    <Form hidden={this.state.qr} onFinish={this.resetTOTP}>
+                        <Form.Item {...formTailLayout}>
+                            <Button type="primary" htmlType="submit">
+                                重置 TOTP
+                            </Button>
+                        </Form.Item>
+                    </Form>
+                    <Form hidden={!this.state.qr} onFinish={this.confirmTOTP}>
+                        <Form.Item {...formItemLayout} label="使用TOTP应用扫码">
+                            <Image
+                                width={200}
+                                src={"data:image/png;base64, " + this.state.qr}
+                            />
+                        </Form.Item>
+                        <Form.Item
+                            {...formItemLayout}
+                            name="totp"
+                            label="TOTP"
+                            rules={[]}
+                        >
+                            <Input placeholder="请输入显示的数字" />
+                        </Form.Item>
+                        <Form.Item {...formTailLayout}>
+                            <Button type="primary" htmlType="submit">
+                                确认
+                            </Button>
+                        </Form.Item>
+                    </Form>
+                </Content>
             </>
         );
     }