add feature# 增加oa登录入口

-from enum import auto
 from logging import error
-from unittest import result
 from flasgger import swag_from
 from app.util import BlueprintApi
-from app.util import BlueprintApi
-from flask import Blueprint, render_template, redirect, request, session, jsonify, flash
+from flask import Blueprint, render_template, redirect, request, session, jsonify, flash, make_response
 from .models import *
 from .oauth2 import authorization, generate_user_info, require_oauth
 from authlib.oauth2 import OAuth2Error
@@ -16,6 +13,9 @@ import time
 from app.models import SM3, AESHelper
 from app.util.component.StructurePrint import StructurePrint
 import traceback
+from oauthlib import oauth2
+import requests
+from app.modules.auth.models import OAuth2Token, User, db
 
 
 def current_user():
@@ -42,6 +42,10 @@ class DataManager(BlueprintApi):
     @bp.route("/authorize", methods=("GET", "POST"))
     def authorize():
         user = current_user()
+        request2 = authorization.create_oauth2_request(request)
+        grant2 = authorization.get_authorization_grant(request=request2)
+        redirect_uri = grant2.validate_authorization_request()
+        session["redirect_uri"] = redirect_uri
         if request.method == "GET":
             # 没有登录，跳转到登录界面
             try:
@@ -212,3 +216,96 @@ class DataManager(BlueprintApi):
             new_pwd = SM3.encode(p)
             result[p] = new_pwd
         return result
+
+    '''
+    三方登录：OA
+    '''
+    @staticmethod
+    @bp.route("/oa", methods=["GET"])
+    def oa_authorization():
+        client = oauth2.WebApplicationClient(
+            configure.OA_OAUTH["client_id"])
+        state = client.state_generator()
+        auth_uri = client.prepare_request_uri(
+            configure.OA_OAUTH["authorization_endpoint"], configure.OA_OAUTH["redirect_uri"], configure.OA_OAUTH["scope"], state)
+        session["oauth_state"] = state
+        return redirect(auth_uri)
+
+    '''
+    oa三方登录回调
+    '''
+    @staticmethod
+    @bp.route("/oa/callback", methods=["GET"])
+    def oa_callback():
+
+        client = oauth2.WebApplicationClient(
+            configure.OA_OAUTH["client_id"])
+
+        # 获取code
+        code = client.parse_request_uri_response(
+            request.url, session["oauth_state"]).get("code")
+
+        if code == None:
+            return "登录失败"
+
+        # 获取token
+        body = client.prepare_request_body(
+            code, redirect_uri=configure.OA_OAUTH["redirect_uri"], client_secret=configure.OA_OAUTH["client_secret"])
+
+        r = requests.post(configure.OA_OAUTH["token_endpoint"], body, headers={
+                          "Content-Type": "application/x-www-form-urlencoded"})
+
+        tokeninfo = r.json()
+        access_token = tokeninfo.get("access_token")
+
+        if access_token:
+            # 获取用户信息
+            userinfo_url = configure.OA_OAUTH["userinfo_endpoint"]
+            user_request = requests.get(userinfo_url, headers={
+                                        "Authorization": "Bearer %s" % access_token})
+            userinfo = user_request.json()
+            user_name = userinfo.get("user_name")
+            display_name = userinfo.get("displayname")
+
+            # 默认关联dmap用户
+            try:
+                user = User.query.filter_by(
+                    username=user_name).first()
+            except error as e:
+                user = None
+
+            # 用户不存在，创建用户
+            if not user:
+                user = User(username=user_name, password=SM3.encode('DMap@123'), role='dataman',
+                            phone='', company='', position='', email='',
+                            create_time=time.strftime(
+                                "%Y-%m-%d %H:%M:%S", time.localtime()),
+                            update_time=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
+                db.session.add(user)
+                db.session.commit()
+
+            # dmap token授权
+            session["id"] = user.id
+            
+            # 存入数据库
+            token = OAuth2Token(
+                client_id=configure.OA_OAUTH["client_id"],
+                token_type=tokeninfo.get("token_type"),
+                access_token=access_token,
+                scope=tokeninfo.get("scope"),
+                expires_in=tokeninfo.get("expires_in"),
+                user_id=user.id
+            )
+            db.session.add(token)
+            db.session.commit()
+
+            redirect_uri = session["redirect_uri"]
+            if not redirect_uri:
+                redirect_uri = '/'
+                
+            response = make_response(redirect(redirect_uri))
+            response.set_cookie('accessToken', access_token,max_age=604_800)
+            
+            return response
+        else:
+            return redirect('/')

@@ -24,6 +24,7 @@ class User (db.Model):
     create_time = Column(Time)
     update_time = Column(Time)
     role = Column(Text)
+    #origin = Column(Text, default="dmap")
 
     def __str__(self):
         return self.username

@@ -76,6 +76,8 @@
             立即登录
           </button>
         </div>
+        <div>——或者——</div>
+        <a href="/auth/oa"> 城信所统一用户认证 </a>
       </form>
     </div>
     <div class="clear"></div>