manyi/library/wechatpayv3/core.py

# -*- coding: utf-8 -*-
# !/usr/bin/env python

import json
import logging
import threading
from datetime import datetime

import requests

from library.wechatbase.exceptions import InvalidSignatureException, APILimitedException, \
    WechatNetworkException, WeChatException
from . import update_certificates
from .type import RequestType, SignType
from .utils import (aes_decrypt, build_authorization, hmac_sign,
                    load_certificate, load_private_key, rsa_decrypt,
                    rsa_encrypt, rsa_sign, rsa_verify)


class Core(object):
    def __init__(self, mchid, cert_serial_no, private_key, apiv3_key, certificate_str_list = None,
                 logger = None, proxy = None, timeout = None, auto_retry = False):
        self._proxy = proxy
        self._mchid = mchid
        self._cert_serial_no = cert_serial_no
        self._private_key = load_private_key(private_key)
        self._apiv3_key = apiv3_key
        self._gate_way = 'https://api.mch.weixin.qq.com'
        self._certificates = []

        self.auto_retry = auto_retry
        self.my_retry_count = threading.local()

        self.timeout = timeout

        self._logger = logger or logging.getLogger(__name__)

        if certificate_str_list:
            for certificate_str in certificate_str_list:
                certificate = load_certificate(certificate_str)
                if not certificate:
                    continue

                now = datetime.utcnow()
                if now < certificate.not_valid_before or now > certificate.not_valid_after:
                    continue

                self._certificates.append(load_certificate(certificate_str))

        if not self._certificates:
            self._update_certificates()

        if not self._certificates:
            raise Exception('No wechatpay platform certificate, please double check your init params.')

    @property
    def mchid(self):
        return self._mchid

    @property
    def retry_count(self):
        if not hasattr(self.my_retry_count, 'retry_count'):
            self.my_retry_count.retry_count = 0
        return self.my_retry_count.retry_count

    @retry_count.setter
    def retry_count(self, retry_count):
        self.my_retry_count.retry_count = retry_count

    def _update_certificates(self):
        path = '/v3/certificates'
        self._certificates = []

        message = self.request(path, skip_verify = True)

        cert_str_list = []

        data = message.get('data')
        for value in data:
            serial_no = value.get('serial_no')
            effective_time = value.get('effective_time')
            expire_time = value.get('expire_time')
            encrypt_certificate = value.get('encrypt_certificate')
            algorithm = nonce = associated_data = ciphertext = None
            if encrypt_certificate:
                algorithm = encrypt_certificate.get('algorithm')
                nonce = encrypt_certificate.get('nonce')
                associated_data = encrypt_certificate.get('associated_data')
                ciphertext = encrypt_certificate.get('ciphertext')
            if not (
                    serial_no and effective_time and expire_time and algorithm and nonce and associated_data and ciphertext):
                continue
            cert_str = aes_decrypt(
                nonce = nonce,
                ciphertext = ciphertext,
                associated_data = associated_data,
                apiv3_key = self._apiv3_key)
            certificate = load_certificate(cert_str)
            if not certificate:
                continue
            now = datetime.utcnow()
            if now < certificate.not_valid_before or now > certificate.not_valid_after:
                continue
            self._certificates.append(certificate)

            cert_str_list.append(cert_str)

        update_certificates.send(self, mchid = self.mchid, cert_str_list = cert_str_list)

    def _verify_signature(self, headers, body):
        signature = headers.get('Wechatpay-Signature')
        timestamp = headers.get('Wechatpay-Timestamp')
        nonce = headers.get('Wechatpay-Nonce')
        serial_no = headers.get('Wechatpay-Serial')
        cert_found = False
        for cert in self._certificates:
            if int('0x' + serial_no, 16) == cert.serial_number:
                cert_found = True
                certificate = cert
                break
        if not cert_found:
            self._update_certificates()
            for cert in self._certificates:
                if int('0x' + serial_no, 16) == cert.serial_number:
                    cert_found = True
                    certificate = cert
                    break
            if not cert_found:
                return False
        if not rsa_verify(timestamp, nonce, body, signature, certificate):
            return False
        return True

    def _decode_result(self, res):
        try:
            result = json.loads(res.content.decode('utf-8', 'ignore'), strict = False)
        except (TypeError, ValueError):
            # Return origin response object if we can not decode it as JSON
            self._logger.debug('Can not decode response as JSON', exc_info = True)
            return res
        return result

    def _handle_error(self, path, res, errcode, errmsg, **kwargs):
        if errcode == 'FREQUENCY_LIMITED':
            raise APILimitedException(
                errCode = errcode,
                errMsg = errmsg,
                client = self,
                request = res.request,
                response = res)

        auto_retry = False
        if errcode in ('SYSTEM_ERROR',):
            auto_retry = True

        if not auto_retry:
            raise WeChatException(
                errCode = errcode,
                errMsg = errmsg,
                client = self,
                request = res.request,
                response = res)
        else:
            self.retry_count = self.retry_count + 1

            if self.retry_count >= 3:
                self._logger.debug('reached the maximum number of retries. url = {}'.format(path))

                self.retry_count = 0
                raise WeChatException(
                    errCode = errcode,
                    errMsg = errmsg,
                    client = self,
                    request = res.request,
                    response = res)
            else:
                self._logger.debug('retry wechat url = {}'.format(path))
                return self.request(path, **kwargs)

    def _handle_result(self, path, res, **kwargs):
        if not isinstance(res, dict):
            result = self._decode_result(res)
        else:
            result = res

        self._logger.debug('Response for Wechat API<url={}> result: {}'.format(path, result))

        if not isinstance(result, dict):
            return result

        if 'code' in result:
            return self._handle_error(path, res, result)

        return result

    def request(self, path, method = RequestType.GET, data = None, skip_verify = False, sign_data = None, files = None,
                cipher_data = False, headers = {}):
        if files:
            headers.update({'Content-Type': 'multipart/form-data'})
        else:
            headers.update({'Content-Type': 'application/json'})

        headers.update({'Accept': 'application/json'})
        headers.update({'User-Agent': 'weifule/v1.0.0 (https://www.washpayer.com)'})

        if cipher_data:
            headers.update({'Wechatpay-Serial': hex(self._last_certificate().serial_number)[2:].upper().rstrip("L")})

        headers.update({'Authorization': build_authorization(
            path, method.value, self._mchid, self._cert_serial_no,
            self._private_key, data = sign_data if sign_data else data)})

        self._logger.debug('Request url: %s' % self._gate_way + path)
        self._logger.debug('Request type: %s' % method.value)
        self._logger.debug('Request headers: %s' % headers)
        self._logger.debug('Request data: %s' % data)

        kwargs = {'headers': headers, 'proxies': self._proxy}

        if method == RequestType.GET:
            pass
        elif method == RequestType.POST:
            kwargs.update({
                'json': None if files else data,
                'data': data if files else None,
                'files': files
            })
        elif method in [RequestType.PATCH, RequestType.PUT]:
            kwargs.update({
                'json': data
            })
        elif method == RequestType.DELETE:
            pass
        else:
            raise Exception('wechatpayv3 does no support this request type.')

        with requests.sessions.Session() as _session:
            res = _session.request(
                method = method.value,
                url = self._gate_way + path,
                timeout = self.timeout,
                **kwargs)

            self._logger.debug('Response status code: %s' % res.status_code)
            self._logger.debug('Response headers: %s' % res.headers)
            self._logger.debug('Response content: %s' % res.text)

            try:
                res.raise_for_status()
            except requests.RequestException as reqe:
                if not isinstance(res, dict):
                    result = self._decode_result(res)
                else:
                    result = res

                if isinstance(result, dict) and 'code' in result:
                    self._handle_error(
                        path, res, result['code'], result['message'], method = method, data = data,
                        skip_verify = skip_verify,
                        sign_data = sign_data, files = files, cipher_data = cipher_data, headers = headers)
                else:
                    raise WechatNetworkException(
                        errCode = 'HTTP{}'.format(res.status_code),
                        errMsg = result or reqe.message,
                        client = self,
                        request = reqe.request,
                        response = reqe.response)

            if not skip_verify:
                if not self._verify_signature(res.headers, res.text):
                    raise InvalidSignatureException()

            if res.status_code == 204:
                return {}
            else:
                if not isinstance(res, dict):
                    result = self._decode_result(res)
                else:
                    result = res

                if not isinstance(result, dict):
                    return result

                if 'code' in result:
                    return self._handle_error(
                        path, res, result['code'], result['message'], method = method, data = data,
                        skip_verify = skip_verify,
                        sign_data = sign_data, files = files, cipher_data = cipher_data, headers = headers)
                else:
                    return result

    def sign(self, data, sign_type = SignType.RSA_SHA256):
        if sign_type == SignType.RSA_SHA256:
            sign_str = '\n'.join(data) + '\n'
            return rsa_sign(self._private_key, sign_str)
        elif sign_type == SignType.HMAC_SHA256:
            key_list = sorted(data.keys())
            sign_str = ''
            for k in key_list:
                v = data[k]
                sign_str += str(k) + '=' + str(v) + '&'
            sign_str += 'key=' + self._apiv3_key
            return hmac_sign(self._apiv3_key, sign_str)
        else:
            raise ValueError('unexpected value of sign_type.')

    def decrypt_callback(self, headers, body):
        if isinstance(body, bytes):
            body = body.decode('UTF-8')
        if self._logger:
            self._logger.debug('Callback Header: %s' % headers)
            self._logger.debug('Callback Body: %s' % body)
        if not self._verify_signature(headers, body):
            return None
        data = json.loads(body)
        resource_type = data.get('resource_type')
        if resource_type != 'encrypt-resource':
            return None
        resource = data.get('resource')
        if not resource:
            return None
        algorithm = resource.get('algorithm')
        if algorithm != 'AEAD_AES_256_GCM':
            raise Exception('wechatpayv3 does not support this algorithm')
        nonce = resource.get('nonce')
        ciphertext = resource.get('ciphertext')
        associated_data = resource.get('associated_data')
        if not (nonce and ciphertext):
            return None
        if not associated_data:
            associated_data = ''
        result = aes_decrypt(
            nonce = nonce,
            ciphertext = ciphertext,
            associated_data = associated_data,
            apiv3_key = self._apiv3_key)
        if self._logger:
            self._logger.debug('Callback resource: %s' % result)
        return result

    def callback(self, headers, body):
        if isinstance(body, bytes):
            body = body.decode('UTF-8')
        result = self.decrypt_callback(headers = headers, body = body)
        if result:
            data = json.loads(body)
            data.update({'resource': json.loads(result)})
            return data
        else:
            return result

    def decrypt(self, ciphtext):
        return rsa_decrypt(ciphertext = ciphtext, private_key = self._private_key)

    def encrypt(self, text):
        return rsa_encrypt(text = text, certificate = self._last_certificate())

    def _last_certificate(self):
        if not self._certificates:
            self._update_certificates()
        certificate = self._certificates[0]
        for cert in self._certificates:
            if certificate.not_valid_after < cert.not_valid_after:
                certificate = cert
        return certificate