base.py 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. # -*- coding: utf-8 -*-
  2. from __future__ import absolute_import, unicode_literals
  3. import struct
  4. import socket
  5. import base64
  6. from library import to_text, to_binary, random_string, byte2int
  7. from library.wechatpy.crypto.pkcs7 import PKCS7Encoder
  8. try:
  9. from library.wechatpy.crypto.cryptography import WeChatCipher
  10. except ImportError:
  11. try:
  12. from library.wechatpy.crypto.pycrypto import WeChatCipher
  13. except ImportError:
  14. raise Exception('You must install either cryptography or pycryptodome!')
  15. class BasePrpCrypto(object):
  16. def __init__(self, key):
  17. self.cipher = WeChatCipher(key)
  18. def get_random_string(self):
  19. return random_string(16)
  20. def _encrypt(self, text, _id):
  21. text = to_binary(text)
  22. tmp_list = []
  23. tmp_list.append(to_binary(self.get_random_string()))
  24. length = struct.pack(b'I', socket.htonl(len(text)))
  25. tmp_list.append(length)
  26. tmp_list.append(text)
  27. tmp_list.append(to_binary(_id))
  28. text = b''.join(tmp_list)
  29. text = PKCS7Encoder.encode(text)
  30. ciphertext = to_binary(self.cipher.encrypt(text))
  31. return base64.b64encode(ciphertext)
  32. def _decrypt(self, text, _id, exception=None):
  33. text = to_binary(text)
  34. plain_text = self.cipher.decrypt(base64.b64decode(text))
  35. padding = byte2int(plain_text[-1])
  36. content = plain_text[16:-padding]
  37. xml_length = socket.ntohl(struct.unpack(b'I', content[:4])[0])
  38. xml_content = to_text(content[4:xml_length + 4])
  39. from_id = to_text(content[xml_length + 4:])
  40. if from_id != _id:
  41. exception = exception or Exception
  42. raise exception()
  43. return xml_content