Completed
Pull Request — master (#138)
by
unknown
20:13
created

BaseWeChatCrypto._encrypt_message()   B

Complexity

Conditions 2

Size

Total Lines 24

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 9
CRAP Score 2.004
Metric Value
cc 2
dl 0
loc 24
ccs 9
cts 10
cp 0.9
crap 2.004
rs 8.9713
1
# -*- coding: utf-8 -*-
2 10
"""
3
    wechatpy.crypto
4
    ~~~~~~~~~~~~~~~~
5
6
    This module provides some crypto tools for WeChat and WeChat enterprise
7
8
    :copyright: (c) 2014 by messense.
9
    :license: MIT, see LICENSE for more details.
10
"""
11 10
from __future__ import absolute_import, unicode_literals
12 10
import time
13 10
import base64
14
15 10
from wechatpy.utils import to_text, to_binary, WeChatSigner
16 10
from wechatpy.exceptions import (
17
    InvalidAppIdException,
18
    InvalidSignatureException
19
)
20 10
from wechatpy.crypto.base import BasePrpCrypto
21
22
23 10
def _get_signature(token, timestamp, nonce, encrypt):
24 10
    signer = WeChatSigner()
25 10
    signer.add_data(token, timestamp, nonce, encrypt)
26 10
    return signer.signature
27
28
29 10
class PrpCrypto(BasePrpCrypto):
30
31 10
    def encrypt(self, text, app_id):
32
        return self._encrypt(text, app_id)
33
34 10
    def decrypt(self, text, app_id):
35
        return self._decrypt(text, app_id, InvalidAppIdException)
36
37
38 10
class BaseWeChatCrypto(object):
39
40 10
    def __init__(self, token, encoding_aes_key, _id):
41 10
        encoding_aes_key = to_binary(encoding_aes_key + '=')
42 10
        self.key = base64.b64decode(encoding_aes_key)
43 10
        assert len(self.key) == 32
44 10
        self.token = token
45 10
        self._id = _id
46
47 10
    def _check_signature(self,
48
                         signature,
49
                         timestamp,
50
                         nonce,
51
                         echo_str,
52
                         crypto_class=None):
53 10
        _signature = _get_signature(self.token, timestamp, nonce, echo_str)
54 10
        if _signature != signature:
55 10
            raise InvalidSignatureException()
56 10
        pc = crypto_class(self.key)
57 10
        return pc.decrypt(echo_str, self._id)
58
59 10
    def _encrypt_message(self,
60
                         msg,
61
                         nonce,
62
                         timestamp=None,
63
                         crypto_class=None):
64 10
        from wechatpy.replies import BaseReply
65
66 10
        xml = """<xml>
67
<Encrypt><![CDATA[{encrypt}]]></Encrypt>
68
<MsgSignature><![CDATA[{signature}]]></MsgSignature>
69
<TimeStamp>{timestamp}</TimeStamp>
70
<Nonce><![CDATA[{nonce}]]></Nonce>
71
</xml>"""
72 10
        if isinstance(msg, BaseReply):
73
            msg = msg.render()
74 10
        timestamp = timestamp or to_binary(int(time.time()))
75 10
        pc = crypto_class(self.key)
76 10
        encrypt = to_text(pc.encrypt(msg, self._id))
77 10
        signature = _get_signature(self.token, timestamp, nonce, encrypt)
78 10
        return to_text(xml.format(
79
            encrypt=encrypt,
80
            signature=signature,
81
            timestamp=timestamp,
82
            nonce=nonce
83
        ))
84
85 10
    def _decrypt_message(self,
86
                         msg,
87
                         signature,
88
                         timestamp,
89
                         nonce,
90
                         crypto_class=None):
91 10
        if not isinstance(msg, dict):
92 10
            import xmltodict
93
94 10
            msg = xmltodict.parse(to_text(msg))['xml']
95
96 10
        encrypt = msg['Encrypt']
97 10
        _signature = _get_signature(self.token, timestamp, nonce, encrypt)
98 10
        if _signature != signature:
99
            raise InvalidSignatureException()
100 10
        pc = crypto_class(self.key)
101 10
        return pc.decrypt(encrypt, self._id)
102
103
104 10
class WeChatCrypto(BaseWeChatCrypto):
105
106 10
    def __init__(self, token, encoding_aes_key, app_id):
107 10
        super(WeChatCrypto, self).__init__(token, encoding_aes_key, app_id)
108 10
        self.app_id = app_id
109
110 10
    def encrypt_message(self, msg, nonce, timestamp=None):
111
        return self._encrypt_message(msg, nonce, timestamp, PrpCrypto)
112
113 10
    def decrypt_message(self, msg, signature, timestamp, nonce):
114
        return self._decrypt_message(
115
            msg,
116
            signature,
117
            timestamp,
118
            nonce,
119
            PrpCrypto
120
        )
121