From 0af4a3dbdc15a1a9314e2bdcbeb68846b067502c Mon Sep 17 00:00:00 2001
From: Richard Chien <richardchienthebest@gmail.com>
Date: Thu, 29 Dec 2016 23:45:34 +0800
Subject: [PATCH] Add weixin support and adapt to mojo-webqq 2.0

---
 README.md                     |  4 +-
 apiclient.py                  | 71 ++++++++++++++++++++++++++++++-----
 app.py                        | 35 +++++++++++++++--
 command.py                    | 36 +++++++++++-------
 commands/core.py              | 14 ++++---
 commands/note.py              | 18 +++++++++
 commands/scheduler.py         | 28 +++++++++++++-
 commands/sudo.py              | 10 +++++
 docker-compose.yml            | 27 ++++++++++---
 docs/Write_Command.md         |  4 ++
 filters/command_dispatcher.py |  8 +---
 filters/message_logger.py     |  1 +
 little_shit.py                | 38 +++++++++++++++----
 13 files changed, 240 insertions(+), 54 deletions(-)
 create mode 100644 commands/sudo.py

diff --git a/README.md b/README.md
index bb60d6c1..12534d1b 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # XiaoKai Bot 小开机器人
 
-基于 [sjdy521/Mojo-Webqq](https://github.com/sjdy521/Mojo-Webqq) 实现的自动处理 QQ 消息的机器人，支持自定义插件。
+基于 [sjdy521/Mojo-Webqq](https://github.com/sjdy521/Mojo-Webqq) 和 [sjdy521/Mojo-Weixin](https://github.com/sjdy521/Mojo-Weixin) 实现的自动处理 QQ 和微信消息的机器人，支持自定义插件。
 
 ## 快速开始
 
@@ -14,7 +14,7 @@
 
 #### 手动运行
 
-首先需要运行 sjdy521/Mojo-Webqq，具体见它的 GitHub 仓库的使用教程。然后运行：
+首先需要运行 sjdy521/Mojo-Webqq 或 sjdy521/Mojo-Webqq，具体见它们的 GitHub 仓库的使用教程。然后运行：
 
 ```sh
 pip install -r requirements.txt
diff --git a/apiclient.py b/apiclient.py
index e239b3e0..18ef727a 100644
--- a/apiclient.py
+++ b/apiclient.py
@@ -4,18 +4,69 @@ import requests
 
 
 class ApiClient:
-    def __init__(self, base_url):
-        self.url = base_url
+    qq_api_url = os.environ.get('QQ_API_URL')
+    wx_api_url = os.environ.get('WX_API_URL')
 
-    def __getattr__(self, item):
-        newclient = ApiClient(self.url + '/' + item)
-        return newclient
-
-    def __call__(self, *args, **kwargs):
+    def send_group_message(self, content: str, ctx_msg: dict):
         try:
-            return requests.get(self.url, params=kwargs)
+            if ctx_msg.get('via') == 'qq' and self.qq_api_url:
+                params = {'content': content}
+                if ctx_msg.get('group_uid'):
+                    params['uid'] = ctx_msg.get('group_uid')
+                elif ctx_msg.get('group_id'):
+                    params['id'] = ctx_msg.get('group_id')
+                return requests.get(self.qq_api_url + '/send_group_message', params=params)
+            elif ctx_msg.get('via') == 'wx' and self.wx_api_url:
+                params = {'content': content}
+                if ctx_msg.get('group_id'):
+                    params['id'] = ctx_msg.get('group_id')
+                return requests.get(self.wx_api_url + '/send_group_message', params=params)
         except requests.exceptions.ConnectionError:
-            return None
+            pass
+        return None
+
+    def send_discuss_message(self, content: str, ctx_msg: dict):
+        try:
+            if ctx_msg.get('via') == 'qq' and self.qq_api_url:
+                params = {'content': content}
+                if ctx_msg.get('discuss_id'):
+                    params['id'] = ctx_msg.get('discuss_id')
+                return requests.get(self.qq_api_url + '/send_discuss_message', params=params)
+        except requests.exceptions.ConnectionError:
+            pass
+        return None
+
+    def send_friend_message(self, content: str, ctx_msg: dict):
+        try:
+            if ctx_msg.get('via') == 'qq' and self.qq_api_url:
+                params = {'content': content}
+                if ctx_msg.get('sender_uid'):
+                    params['uid'] = ctx_msg.get('sender_uid')
+                elif ctx_msg.get('sender_id'):
+                    params['id'] = ctx_msg.get('sender_id')
+                return requests.get(self.qq_api_url + '/send_friend_message', params=params)
+            elif ctx_msg.get('via') == 'wx' and self.wx_api_url:
+                params = {'content': content}
+                if ctx_msg.get('sender_account'):
+                    params['account'] = ctx_msg.get('sender_account')
+                elif ctx_msg.get('sender_id'):
+                    params['id'] = ctx_msg.get('sender_id')
+                return requests.get(self.wx_api_url + '/send_friend_message', params=params)
+        except requests.exceptions.ConnectionError:
+            pass
+        return None
+
+    def get_group_info(self, via):
+        url = None
+        if via == 'qq':
+            url = self.qq_api_url
+        elif via == 'wx':
+            url = self.wx_api_url
+        if url:
+            try:
+                return requests.get(url + '/get_group_info')
+            except requests.exceptions.ConnectionError:
+                return None
 
 
-client = ApiClient(os.environ.get('QQ_API_URL'))
+client = ApiClient()
diff --git a/app.py b/app.py
index 8e7a90a7..f62c544e 100644
--- a/app.py
+++ b/app.py
@@ -9,14 +9,28 @@ from filter import apply_filters
 app = Flask(__name__)
 
 
-@app.route('/', methods=['POST'])
-def _index():
+@app.route('/qq/', methods=['POST'])
+def _handle_qq_message():
     ctx_msg = request.json
+    ctx_msg['via'] = 'qq'
+    return _main(ctx_msg)
+
+
+@app.route('/wx/', methods=['POST'])
+def _handle_wx_message():
+    ctx_msg = request.json
+    ctx_msg['via'] = 'wx'
+    return _main(ctx_msg)
+
+
+def _main(ctx_msg: dict):
+    _preprocess_ctx_msg(ctx_msg)
     try:
-        if ctx_msg.get('msg_class') != 'recv':
+        if ctx_msg.get('post_type') != 'receive_message':
             raise SkipException
         if not apply_filters(ctx_msg):
             raise SkipException
+        print(ctx_msg)
     except SkipException:
         # Skip this message
         pass
@@ -24,6 +38,21 @@ def _index():
     return '', 204
 
 
+def _preprocess_ctx_msg(ctx_msg: dict):
+    if 'group_uid' in ctx_msg:
+        ctx_msg['group_uid'] = str(ctx_msg['group_uid'])
+    if 'sender_uid' in ctx_msg:
+        ctx_msg['sender_uid'] = str(ctx_msg['sender_uid'])
+    if 'sender_id' in ctx_msg:
+        ctx_msg['sender_id'] = str(ctx_msg['sender_id'])
+    if 'discuss_id' in ctx_msg:
+        ctx_msg['discuss_id'] = str(ctx_msg['discuss_id'])
+    if 'group_id' in ctx_msg:
+        ctx_msg['group_id'] = str(ctx_msg['group_id'])
+    if 'id' in ctx_msg:
+        ctx_msg['id'] = str(ctx_msg['id'])
+
+
 def _load_filters():
     filter_mod_files = filter(
         lambda filename: filename.endswith('.py') and not filename.startswith('_'),
diff --git a/command.py b/command.py
index bf855e38..df391e29 100644
--- a/command.py
+++ b/command.py
@@ -65,7 +65,7 @@ class CommandRegistry:
     # noinspection PyMethodMayBeStatic
     def restrict(self, full_command_only=False, superuser_only=False,
                  group_owner_only=False, group_admin_only=False,
-                 allow_private=True, allow_group=True):
+                 allow_private=True, allow_discuss=True, allow_group=True):
         """
         Give a command some restriction.
         This decorator must be put below all register() decorators.
@@ -81,6 +81,7 @@ class CommandRegistry:
         :param group_owner_only: group owner only when processing group message
         :param group_admin_only: group admin only when processing group message
         :param allow_private: allow private message
+        :param allow_discuss: allow discuss message
         :param allow_group: allow group message
         """
 
@@ -94,6 +95,7 @@ class CommandRegistry:
             func.group_admin_only = group_admin_only
             # Scope
             func.allow_private = allow_private
+            func.allow_discuss = allow_discuss
             func.allow_group = allow_group
             return func
 
@@ -114,9 +116,14 @@ class CommandRegistry:
         if command_name in self.command_map:
             func = self.command_map[command_name]
             if not self._check_scope(func, ctx_msg):
-                raise CommandScopeError(
-                    '群组消息' if ctx_msg.get('type') == 'group_message' else '私聊消息'
-                )
+                msg_type = ctx_msg.get('type')
+                if msg_type == 'group_message':
+                    msg_type_str = '群组消息'
+                elif msg_type == 'discuss_message':
+                    msg_type_str = '讨论组消息'
+                else:
+                    msg_type_str = '私聊消息'
+                raise CommandScopeError(msg_type_str)
             if not self._check_permission(func, ctx_msg):
                 raise CommandPermissionError
             return func(args_text, ctx_msg, **options)
@@ -133,8 +140,10 @@ class CommandRegistry:
         allowed_msg_type = set()
         if func.allow_group:
             allowed_msg_type.add('group_message')
+        if func.allow_discuss:
+            allowed_msg_type.add('discuss_message')
         if func.allow_private:
-            allowed_msg_type.add('message')
+            allowed_msg_type.add('friend_message')
 
         if ctx_msg.get('type') in allowed_msg_type:
             return True
@@ -156,34 +165,35 @@ class CommandRegistry:
 
         try:
             if func.superuser_only:
-                check(str(ctx_msg.get('sender_qq')) == os.environ.get('SUPER_USER_QQ'))
-            if ctx_msg.get('type') == 'group_message':
+                raise SkipException
+            if ctx_msg.get('type') == 'group_message' and ctx_msg.get('via') == 'qq':
                 allowed_roles = {'owner', 'admin', 'member'}
                 if func.group_admin_only:
                     allowed_roles = allowed_roles.intersection({'owner', 'admin'})
                 if func.group_owner_only:
                     allowed_roles = allowed_roles.intersection({'owner'})
                 groups = list(filter(
-                    lambda g: str(g.get('gnumber')) == str(ctx_msg.get('gnumber')),
-                    api.get_group_info().json()
+                    lambda g: g.get('group_uid') == ctx_msg.get('group_uid'),
+                    api.get_group_info(via='qq').json()
                 ))
                 if len(groups) <= 0 or 'member' not in groups[0]:
                     # This is strange, not likely happens
                     raise SkipException
 
                 members = list(filter(
-                    lambda m: str(m.get('qq')) == str(ctx_msg.get('sender_qq')),
+                    lambda m: str(m.get('uid')) == str(ctx_msg.get('sender_uid')),
                     groups[0].get('member')
                 ))
                 if len(members) <= 0 or members[0].get('role') not in allowed_roles:
                     # This is strange, not likely happens
                     raise SkipException
         except SkipException:
-            if not str(ctx_msg.get('sender_qq')) == os.environ.get('SUPER_USER_QQ'):
-                # Not allowed
+            if ctx_msg.get('via') == 'qq' and ctx_msg.get('sender_uid') != os.environ.get('QQ_SUPER_USER'):
+                return False
+            elif ctx_msg.get('via') == 'wx' and ctx_msg.get('sender_account') != os.environ.get('WX_SUPER_USER'):
                 return False
 
-        # Still alive, so let it go
+        # Still alive, let go
         return True
 
     def has(self, command_name):
diff --git a/commands/core.py b/commands/core.py
index 7ce33b1f..e5800ac6 100644
--- a/commands/core.py
+++ b/commands/core.py
@@ -11,9 +11,11 @@ __registry__ = cr = CommandRegistry()
 def echo(args_text, ctx_msg):
     msg_type = ctx_msg.get('type')
     if msg_type == 'group_message':
-        api.send_group_message(gnumber=ctx_msg.get('gnumber'), content=args_text)
-    elif msg_type == 'message':
-        api.send_message(qq=ctx_msg.get('sender_qq'), content=args_text)
+        api.send_group_message(content=args_text, ctx_msg=ctx_msg)
+    elif msg_type == 'discuss_message':
+        api.send_discuss_message(content=args_text, ctx_msg=ctx_msg)
+    elif msg_type == 'friend_message':
+        api.send_friend_message(content=args_text, ctx_msg=ctx_msg)
 
 
 @cr.register('chat', '聊天')
@@ -23,8 +25,10 @@ def chat(args_text, ctx_msg):
         'key': os.environ.get('TURING123_API_KEY'),
         'info': args_text
     }
-    if 'sender_qq' in ctx_msg:
-        data['userid'] = ctx_msg.get('sender_qq')
+    if ctx_msg.get('sender_uid'):
+        data['userid'] = ctx_msg.get('sender_uid')
+    elif ctx_msg.get('sender_id'):
+        data['userid'] = ctx_msg.get('sender_id')[-32:]
     resp = requests.post(url, data=data)
     if resp.status_code == 200:
         json = resp.json()
diff --git a/commands/note.py b/commands/note.py
index a90e6f8a..e8da5251 100644
--- a/commands/note.py
+++ b/commands/note.py
@@ -1,4 +1,5 @@
 import sqlite3
+from functools import wraps
 from datetime import datetime
 
 import pytz
@@ -25,6 +26,19 @@ def _open_db_conn():
     return conn
 
 
+def _check_target(func):
+    @wraps(func)
+    def wrapper(args_text, ctx_msg, *args, **kwargs):
+        target = get_target(ctx_msg)
+        if not target:
+            core.echo('似乎出错了，请稍后再试吧～', ctx_msg)
+            return
+        else:
+            return func(args_text, ctx_msg, *args, **kwargs)
+
+    return wrapper
+
+
 _cmd_take = 'note.take'
 _cmd_remove = 'note.remove'
 
@@ -32,6 +46,7 @@ _cmd_remove = 'note.remove'
 @cr.register('记笔记', '添加笔记')
 @cr.register('take', 'add', hidden=True)
 @cr.restrict(group_admin_only=True)
+@_check_target
 def take(args_text, ctx_msg, allow_interactive=True):
     source = get_source(ctx_msg)
     if allow_interactive and (not args_text or has_session(source, _cmd_take)):
@@ -52,6 +67,7 @@ def take(args_text, ctx_msg, allow_interactive=True):
 
 @cr.register('列出所有笔记')
 @cr.register('list', hidden=True)
+@_check_target
 def list_all(_, ctx_msg):
     conn = _open_db_conn()
     target = get_target(ctx_msg)
@@ -74,6 +90,7 @@ def list_all(_, ctx_msg):
 @cr.register('删除笔记')
 @cr.register('remove', 'delete', hidden=True)
 @cr.restrict(group_admin_only=True)
+@_check_target
 def remove(args_text, ctx_msg, allow_interactive=True):
     source = get_source(ctx_msg)
     if allow_interactive and (not args_text or has_session(source, _cmd_remove)):
@@ -101,6 +118,7 @@ def remove(args_text, ctx_msg, allow_interactive=True):
 @cr.register('清空笔记', '清空所有笔记', '删除所有笔记')
 @cr.register('clear', hidden=True)
 @cr.restrict(group_admin_only=True)
+@_check_target
 def clear(_, ctx_msg):
     conn = _open_db_conn()
     target = get_target(ctx_msg)
diff --git a/commands/scheduler.py b/commands/scheduler.py
index c4f9d26e..7006e3b3 100644
--- a/commands/scheduler.py
+++ b/commands/scheduler.py
@@ -1,6 +1,6 @@
 import os
 import re
-from functools import reduce
+from functools import reduce, wraps
 
 import pytz
 from apscheduler.schedulers.background import BackgroundScheduler
@@ -60,8 +60,22 @@ def _call_commands(job_id, command_list, ctx_msg):
             )
 
 
+def _check_target(func):
+    @wraps(func)
+    def wrapper(args_text, ctx_msg, internal=False, *args, **kwargs):
+        target = get_target(ctx_msg)
+        if not target:
+            _send_fail_to_get_target_msg(ctx_msg, internal)
+            return None
+        else:
+            return func(args_text, ctx_msg, internal, *args, **kwargs)
+
+    return wrapper
+
+
 @cr.register('add_job', 'add-job', 'add')
 @cr.restrict(full_command_only=True, group_admin_only=True)
+@_check_target
 def add_job(args_text, ctx_msg, internal=False):
     if args_text.strip() in ('', 'help', '-h', '--help') and not internal:
         _send_add_job_help_msg(ctx_msg, internal)
@@ -150,6 +164,7 @@ def add_job(args_text, ctx_msg, internal=False):
 
 @cr.register('remove_job', 'remove-job', 'remove')
 @cr.restrict(full_command_only=True, group_admin_only=True)
+@_check_target
 def remove_job(args_text, ctx_msg, internal=False):
     job_id_without_suffix = args_text.strip()
     if not job_id_without_suffix:
@@ -167,6 +182,7 @@ def remove_job(args_text, ctx_msg, internal=False):
 
 @cr.register('get_job', 'get-job', 'get')
 @cr.restrict(full_command_only=True)
+@_check_target
 def get_job(args_text, ctx_msg, internal=False):
     job_id_without_suffix = args_text.strip()
     if not job_id_without_suffix:
@@ -190,6 +206,7 @@ def get_job(args_text, ctx_msg, internal=False):
 
 @cr.register('list_jobs', 'list-jobs', 'list')
 @cr.restrict(full_command_only=True)
+@_check_target
 def list_jobs(_, ctx_msg, internal=False):
     target = get_target(ctx_msg)
     job_id_suffix = '_' + target
@@ -219,6 +236,15 @@ def _send_text(text, ctx_msg, internal):
         core.echo(text, ctx_msg)
 
 
+def _send_fail_to_get_target_msg(ctx_msg, internal):
+    _send_text(
+        '无法获取 target，可能因为不支持当前消息类型（如，不支持微信群组消息）'
+        '或由于延迟还没能加载到用户的固定 ID（如，微信号）',
+        ctx_msg,
+        internal
+    )
+
+
 def _send_add_job_help_msg(ctx_msg, internal):
     _send_text(
         '此为高级命令！如果你不知道自己在做什么，请不要使用此命令。\n\n'
diff --git a/commands/sudo.py b/commands/sudo.py
new file mode 100644
index 00000000..9d01e691
--- /dev/null
+++ b/commands/sudo.py
@@ -0,0 +1,10 @@
+from command import CommandRegistry
+from commands import core
+
+__registry__ = cr = CommandRegistry()
+
+
+@cr.register('test')
+@cr.restrict(full_command_only=True, superuser_only=True)
+def test(_, ctx_msg):
+    core.echo('Your are the superuser!', ctx_msg)
diff --git a/docker-compose.yml b/docker-compose.yml
index 38edfddd..178524a4 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,8 +1,8 @@
 version: '2'
 services:
-  mojo-smartqq-api:
-    image: sjdy521/mojo-webqq
-    container_name: mojo-smartqq-api
+  mojo-webqq-api:
+    image: daocloud.io/richardchien/mojo-webqq
+    container_name: mojo-webqq-api
     networks:
       - my-net
     expose:
@@ -11,7 +11,20 @@ services:
       - /tmp:/tmp
     environment:
       - PORT=5000
-      - POST_API=http://qq-bot:8888
+      - POST_API=http://qq-bot:8888/qq/
+    restart: always
+  mojo-weixin-api:
+    image: daocloud.io/richardchien/mojo-weixin
+    container_name: mojo-weixin-api
+    networks:
+      - my-net
+    expose:
+      - '5001'
+    volumes:
+      - /tmp:/tmp
+    environment:
+      - PORT=5001
+      - POST_API=http://qq-bot:8888/wx/
     restart: always
   qq-bot:
     image: daocloud.io/richardchien/xiaokai-bot
@@ -26,8 +39,10 @@ services:
       - TURING123_API_KEY=YOUR_API_KEY
       - HOST=0.0.0.0
       - PORT=8888
-      - QQ_API_URL=http://mojo-smartqq-api:5000/openqq
-      - SUPER_USER_QQ=12345678
+      - QQ_API_URL=http://mojo-webqq-api:5000/openqq
+      - QQ_SUPER_USER=12345678
+      - WX_API_URL=http://mojo-weixin-api:5001/openwx
+      - WX_SUPER_USER=abcdedfh
     restart: always
 networks:
   my-net:
diff --git a/docs/Write_Command.md b/docs/Write_Command.md
index 60113b8a..dab6014d 100644
--- a/docs/Write_Command.md
+++ b/docs/Write_Command.md
@@ -35,6 +35,8 @@ def list_all(args_text, ctx_msg):
 
 这样可以在保持高级用户可以通过简洁的方式调用命令的同时，避免不同命令仓库下同名命令都被调用的问题（因为在默认情况下命令中心在调用命令时，不同仓库中的同名命令都会被依次调用）。
 
+16.12.29 注：由于微信限制，无法获取到群组中成员的身份（普通成员还是管理员或群主），因此这里的对群组的限制在微信上不起效果，超级用户限制在能够获取到发送者微信 ID 的情况下有效。
+
 ## 命令中心 Command Hub
 
 程序启动时加载的命令仓库全部被集中在了命令中心，以 `.py` 文件名（除去后缀）（也即模块名）为仓库名。命令中心实际上应作为单例使用（插件编写者不应当自己创建实例），即 `command.py` 中的 `hub` 对象，类型是 CommandHub，调用它的 `call` 方法将会执行相应的命令，如果调用失败（如命令不存在、没有权限等）会抛出相应的异常。
@@ -73,6 +75,8 @@ Source 表示命令的来源（由谁发出），Target 表示命令将对谁产
 
 至于如何获取 Source 和 Target，可用 `little_shit.py` 中的 `get_source` 和 `get_target` 函数，当然，如果你对默认的行为感到不满意，也可以自己去实现不一样的区分方法。
 
+16.12.29 注：在支持了微信之后，此处有所变化，由于微信消息的限制，有时候无法获得发送者的微信 ID，而群组甚至没有一个固定 ID，因此，这里对 Source 和 Target 做一个精确定义：Source 是一个用来表示当前消息的发送者的唯一值，但重新登录后可能变化，并且每次获取，一定可以获取到；Target 是一个用来表示当前消息所产生的效果需要作用的对象，这个值是永久（或至少长期）不变的，如果当前的消息语境下不存在这样的值，则为 None。
+
 ## 交互式命令 Interactive Command
 
 通常我们会希望一条命令直接在一条消息中发出然后直接执行就好，但这在某些情况下对普通用户并不友好，因此这里支持了交互式命令，也就是在一条命令调用之后，进行多次后续互动，来引导用户完成数据的输入。
diff --git a/filters/command_dispatcher.py b/filters/command_dispatcher.py
index 7a28a5e2..3e5bca61 100644
--- a/filters/command_dispatcher.py
+++ b/filters/command_dispatcher.py
@@ -8,7 +8,6 @@ from filter import add_filter
 from command import CommandNotExistsError, CommandScopeError, CommandPermissionError
 from little_shit import *
 from commands import core
-from apiclient import client as api
 from command import hub as cmdhub
 
 _fallback_command = config.get('fallback_command')
@@ -40,14 +39,11 @@ def _dispatch_command(ctx_msg):
                 raise SkipException
             at_me = '@' + my_group_nick
             if not content.startswith(at_me):
-                my_nick = api.get_user_info().json().get('nick', my_group_nick)
-                at_me = '@' + my_nick
-                if not content.startswith(at_me):
-                    raise SkipException
+                raise SkipException
             content = content[len(at_me):]
         else:
             # Not starts with '@'
-            if ctx_msg.get('type') == 'group_message':
+            if ctx_msg.get('type') == 'group_message' or ctx_msg.get('type') == 'discuss_message':
                 # And it's a group message, so we don't reply
                 raise SkipException
         content = content.lstrip()
diff --git a/filters/message_logger.py b/filters/message_logger.py
index f626d557..6ad58690 100644
--- a/filters/message_logger.py
+++ b/filters/message_logger.py
@@ -4,6 +4,7 @@ from filter import add_filter
 def _log_message(ctx_msg):
     print(ctx_msg.get('sender', '')
           + (('@' + ctx_msg.get('group')) if ctx_msg.get('type') == 'group_message' else '')
+          + (('@' + ctx_msg.get('discuss')) if ctx_msg.get('type') == 'discuss_message' else '')
           + ': ' + ctx_msg.get('content'))
 
 
diff --git a/little_shit.py b/little_shit.py
index 3661d888..e7319d6e 100644
--- a/little_shit.py
+++ b/little_shit.py
@@ -39,21 +39,43 @@ def get_tmp_dir():
 def get_source(ctx_msg):
     """
     Source is used to distinguish the interactive sessions.
+    Note: This value may change after restarting the bot.
+
+    :return: an unique value representing a source, or None if the 'via' field is not recognized
     """
-    if ctx_msg.get('type') == 'group_message':
-        return 'g' + str(ctx_msg.get('gnumber')) + 'p' + str(ctx_msg.get('sender_qq'))
-    else:
-        return 'p' + str(ctx_msg.get('sender_qq'))
+    if ctx_msg.get('via') == 'qq':
+        if ctx_msg.get('type') == 'group_message':
+            return 'g' + ctx_msg.get('group_uid') + 'p' + ctx_msg.get('sender_uid')
+        elif ctx_msg.get('type') == 'discuss_message':
+            return 'd' + ctx_msg.get('discuss_id') + 'p' + ctx_msg.get('sender_uid')
+        else:
+            return 'p' + str(ctx_msg.get('sender_uid'))
+    elif ctx_msg.get('via') == 'wx':
+        if ctx_msg.get('type') == 'group_message':
+            return 'g' + ctx_msg.get('group_id') + 'p' + ctx_msg.get('sender_id')
+        else:
+            return 'p' + ctx_msg.get('sender_id')
 
 
 def get_target(ctx_msg):
     """
     Target is used to distinguish the records in database.
+    Note: This value will not change after restarting the bot.
+
+    :return: an unique value representing a target, or None if there is no persistent unique value
     """
-    if ctx_msg.get('type') == 'group_message':
-        return 'g' + str(ctx_msg.get('gnumber'))
-    else:
-        return 'p' + str(ctx_msg.get('sender_qq'))
+    if ctx_msg.get('via') == 'qq':
+        if ctx_msg.get('type') == 'group_message':
+            return 'g' + str(ctx_msg.get('group_uid'))
+        elif ctx_msg.get('type') == 'discuss_message':
+            # TODO: 看看讨论组 ID 重新启动会不会变
+            pass
+        elif ctx_msg.get('type') == 'friend_message':
+            return 'p' + str(ctx_msg.get('sender_uid'))
+    elif ctx_msg.get('via') == 'wx':
+        if ctx_msg.get('type') == 'friend_message' and ctx_msg.get('sender_account'):
+            return 'p' + str(ctx_msg.get('sender_account'))
+    return None
 
 
 def get_command_start_flags():