ADD:添加微信用户验证,一号一用户

wechat_action/login_ad.py

@@ -16,6 +16,7 @@ from wechat_action.sql_models import DB
 import requests
 import json
 import logging
+import re
 
 
 class LogIn:
@@ -71,16 +72,22 @@ class LogIn:
         log_ad.driver.switch_to.window(log_ad.driver.window_handles[0])
         WebDriverWait(log_ad.driver, 100).until(lambda driver: driver.find_element_by_class_name(
             'ui-mr-medium'))
-        # cookie 进行数据库保存
-        wechat_cookies = log_ad.wechat_cookie_pickle()
-        update_res = log_ad.wechat_cookies_table.update() \
-            .where(log_ad.wechat_cookies_table.c.user_id == log_ad.user_id) \
-            .values(cookies=wechat_cookies,
-                    scan_action='done')
+        wechat_cookies, wechat_id_tmp = log_ad.wechat_cookie_pickle()
+
         sql_session = log_ad.db.DBSession()
-        cursor = sql_session.execute(update_res)
-        sql_session.commit()
-        if cursor.rowcount == 0:
+        wechat_id = sql_tools.get_wechat_id_from_cookies(log_ad.user_id, sql_session)
+        if wechat_id:
+            if wechat_id_tmp != wechat_id:
+                raise ValueError("微信账号,非以前老账号,登录失败")
+
+            # cookie 进行数据库保存
+            update_res = log_ad.wechat_cookies_table.update() \
+                .where(log_ad.wechat_cookies_table.c.user_id == log_ad.user_id) \
+                .values(cookies=wechat_cookies,
+                        scan_action='done')
+            sql_session.execute(update_res)
+            sql_session.commit()
+        else:
             wechat_cookies_info = {'user_id': log_ad.user_id, 'cookies': wechat_cookies, 'scan_action': 'done'}
             wechat_insert = sql_tools.save_wechat_cookies(wechat_cookies_info=wechat_cookies_info,
                                                           table_wechat_cookies=log_ad.wechat_cookies_table)
@@ -135,13 +142,25 @@ class LogIn:
         return cookie_dict
 
     def wechat_cookie_pickle(self):
+        wechat_id = None
         self.driver.get('https://a.weixin.qq.com/client')
         WebDriverWait(self.driver, 100).until(
             lambda x: [True for _ in self.driver.get_cookies() if 'token_ticket' == _['name']]
         )
+        WebDriverWait(self.driver, 100).until(
+            lambda x: re.findall('g_tk=(\d+)&', self.driver.page_source)
+        )
         cookies = self.driver.get_cookies()
         cookies_obj = pickle.dumps(cookies)
-        return cookies_obj
+        g_tk = re.findall('g_tk=(\d+)&', self.driver.page_source)[0]
+        wechat_cookies = self.get_cookie(self.driver)
+        wechat_url = 'https://a.weixin.qq.com/cgi-bin/agency/check_login?g_tk={g_tk}&_={time_p}'.format(g_tk=g_tk,
+                                                                                                        time_p=int(
+                                                                                                            time.time() * 1000))
+        rsp = requests.get(url=wechat_url, cookies=wechat_cookies)
+        wechat_id = rsp.json()['data'][0]['wx_id']
+
+        return cookies_obj, wechat_id
 
     def wechat_cookies_check_alive(self, driver_cookies):
         # wechat 检查cookies 是否可用
@@ -165,10 +184,6 @@ class LogIn:
         else:
             return True
 
-    def clock_wechat_user(self):
-        # TODO:用户登录之后,微信号和本地账号进行一个锁死
-        pass
-
     def upadte_user_info(self):
         # TODO: 更新 用户相关信息
         #  每次登录就更新一次相关数据,------公众号相关数据,人群报相关数据

wechat_action/sql_tools.py

@@ -112,7 +112,7 @@ def get_scan_action_status(user_id, sql_session):
         diff_time, scan_action = lines[0]
         if scan_action == 'done':
             return 'done'
-        elif scan_action =='doing':
+        elif scan_action == 'doing':
             if diff_time > 1:
                 return 'error'
             else:
@@ -290,6 +290,19 @@ def update_user_scan_action(user_id, sql_session):
     sql_session.commit()
 
 
+def get_wechat_id_from_cookies(user_id, sql_session):
+    sql = '''
+    select wechat_id from wechat_cookies wc 
+    where user_id ='{}'
+    '''.format(user_id)
+    cursor = sql_session.execute(sql)
+    lines = cursor.fetchall()
+    wechat_id = None
+    if lines:
+        wechat_id = lines[0][0]
+    return wechat_id
+
+
 def get_ad_task(sql_session, user_id):
     sql = '''
     select task_name,status,count(*),min(ar.create_time),apt.typesetting from action_record ar