add MFA

spug_api/apps/account/models.py

@@ -19,6 +19,7 @@ class User(models.Model, ModelMixin):
     last_login = models.CharField(max_length=20)
     last_ip = models.CharField(max_length=50)
     role = models.ForeignKey('Role', on_delete=models.PROTECT, null=True)
+    wx_token = models.CharField(max_length=50, null=True)
 
     created_at = models.CharField(max_length=20, default=human_datetime)
     created_by = models.ForeignKey('User', models.PROTECT, related_name='+', null=True)

spug_api/apps/account/views.py

@@ -5,7 +5,8 @@ from django.core.cache import cache
 from django.views.generic import View
 from django.db.models import F
 from libs import JsonParser, Argument, human_datetime, json_response
-from libs.utils import get_request_real_ip
+from libs.utils import get_request_real_ip, generate_random_str
+from libs.spug import send_login_wx_code
 from apps.account.models import User, Role, History
 from apps.setting.utils import AppSetting
 from libs.ldap import LDAP
@@ -18,7 +19,7 @@ import json
 class UserView(View):
     def get(self, request):
         users = []
-        for u in User.objects.filter(is_supper=False, deleted_by_id__isnull=True).annotate(role_name=F('role__name')):
+        for u in User.objects.filter(deleted_by_id__isnull=True).annotate(role_name=F('role__name')):
             tmp = u.to_dict(excludes=('access_token', 'password_hash'))
             tmp['role_name'] = u.role_name
             users.append(tmp)
@@ -30,6 +31,7 @@ class UserView(View):
             Argument('password', help='请输入密码'),
             Argument('nickname', help='请输入姓名'),
             Argument('role_id', type=int, help='请选择角色'),
+            Argument('wx_token', required=False),
         ).parse(request.body)
         if error is None:
             if User.objects.filter(username=form.username, deleted_by_id__isnull=True).exists():
@@ -46,6 +48,7 @@ class UserView(View):
             Argument('password', required=False),
             Argument('nickname', required=False),
             Argument('role_id', required=False),
+            Argument('wx_token', required=False),
             Argument('is_active', type=bool, required=False),
         ).parse(request.body, True)
         if error is None:
@@ -156,10 +159,10 @@ def login(request):
     form, error = JsonParser(
         Argument('username', help='请输入用户名'),
         Argument('password', help='请输入密码'),
+        Argument('captcha', required=False),
         Argument('type', required=False)
     ).parse(request.body)
     if error is None:
-        x_real_ip = get_request_real_ip(request.headers)
         user = User.objects.filter(username=form.username, type=form.type).first()
         if user and not user.is_active:
             return json_response(error="账户已被系统禁用")
@@ -171,13 +174,13 @@ def login(request):
             if is_success:
                 if not user:
                     user = User.objects.create(username=form.username, nickname=form.username, type=form.type)
-                return handle_user_info(user, x_real_ip)
+                return handle_user_info(request, user, form.captcha)
             elif message:
                 return json_response(error=message)
         else:
             if user and user.deleted_by is None:
                 if user.verify_password(form.password):
-                    return handle_user_info(user, x_real_ip)
+                    return handle_user_info(request, user, form.captcha)
 
         value = cache.get_or_set(form.username, 0, 86400)
         if value >= 3:
@@ -190,8 +193,29 @@ def login(request):
     return json_response(error=error)
 
 
-def handle_user_info(user, x_real_ip):
+def handle_user_info(request, user, captcha):
     cache.delete(user.username)
+    key = f'{user.username}:code'
+    if captcha:
+        code = cache.get(key)
+        if not code:
+            return json_response(error='验证码已失效，请重新获取')
+        if code != captcha:
+            ttl = cache.ttl(key)
+            cache.expire(key, ttl - 100)
+            return json_response(error='验证码错误')
+        cache.delete(key)
+    else:
+        mfa = AppSetting.get_default('MFA', {'enable': False})
+        if mfa['enable']:
+            if not user.wx_token:
+                return json_response(error='已启用登录双重认证，但您的账户未配置微信Token，请联系管理员')
+            code = generate_random_str(6)
+            send_login_wx_code(user.wx_token, code)
+            cache.set(key, code, 300)
+            return json_response({'required_mfa': True})
+
+    x_real_ip = get_request_real_ip(request.headers)
     token_isvalid = user.access_token and len(user.access_token) == 32 and user.token_expired >= time.time()
     user.access_token = user.access_token if token_isvalid else uuid.uuid4().hex
     user.token_expired = time.time() + 8 * 60 * 60

spug_api/libs/spug.py

@@ -13,6 +13,16 @@ spug_server = 'https://api.spug.cc'
 notify_source = 'monitor'
 
 
+def send_login_wx_code(wx_token, code):
+    url = f'{spug_server}/apis/login/wx/'
+    res = requests.post(url, json={'token': wx_token, 'code': code}, timeout=30)
+    if res.status_code != 200:
+        raise Exception(f'status code: {res.status_code}')
+    res = res.json()
+    if res.get('error'):
+        raise Exception(res['error'])
+
+
 class Notification:
     def __init__(self, grp, event, target, title, message, duration):
         self.event = event

spug_web/src/pages/login/index.js

@@ -21,6 +21,8 @@ export default function () {
   const [counter, setCounter] = useState(0);
   const [loading, setLoading] = useState(false);
   const [loginType, setLoginType] = useState('default');
+  const [codeVisible, setCodeVisible] = useState(false);
+  const [codeLoading, setCodeLoading] = useState(false);
 
   useEffect(() => {
     envStore.records = [];
@@ -45,7 +47,11 @@ export default function () {
     formData['type'] = loginType;
     http.post('/api/account/login/', formData)
       .then(data => {
-        if (!data['has_real_ip']) {
+        if (data['required_mfa']) {
+          setCodeVisible(true);
+          setCounter(30);
+          setLoading(false)
+        } else if (!data['has_real_ip']) {
           Modal.warning({
             title: '安全警告',
             className: styles.tips,
@@ -78,7 +84,12 @@ export default function () {
   }
 
   function handleCaptcha() {
-    setCounter(60);
+    setCodeLoading(true);
+    const formData = form.getFieldsValue(['username', 'password']);
+    formData['type'] = loginType;
+    http.post('/api/account/login/', formData)
+      .then(() => setCounter(30))
+      .finally(() => setCodeLoading(false))
   }
 
   return (
@@ -109,7 +120,7 @@ export default function () {
               onPressEnter={handleSubmit}
               prefix={<LockOutlined className={styles.icon}/>}/>
           </Form.Item>
-          <Form.Item name="captcha" className={styles.formItem}>
+          <Form.Item hidden={!codeVisible} name="captcha" className={styles.formItem}>
             <div style={{display: 'flex'}}>
               <Form.Item noStyle name="captcha">
                 <Input
@@ -121,7 +132,8 @@ export default function () {
               {counter > 0 ? (
                 <Button disabled size="large" style={{marginLeft: 8}}>{counter} 秒后重新获取</Button>
               ) : (
-                <Button size="large" style={{marginLeft: 8}} onClick={handleCaptcha}>获取验证码</Button>
+                <Button size="large" loading={codeLoading} style={{marginLeft: 8}}
+                        onClick={handleCaptcha}>获取验证码</Button>
               )}
             </div>
           </Form.Item>

spug_web/src/pages/system/account/Form.js

@@ -71,6 +71,9 @@ export default observer(function () {
             <Link to="/system/role">新建角色</Link>
           </Form.Item>
         </Form.Item>
+        <Form.Item name="wx_token" label="微信Token">
+          <Input placeholder="请输入微信Token"/>
+        </Form.Item>
       </Form>
     </Modal>
   )