diff --git a/itchat/components/messages.py b/itchat/components/messages.py index 2cc1dadd..fff5e302 100644 --- a/itchat/components/messages.py +++ b/itchat/components/messages.py @@ -276,19 +276,38 @@ def send_msg(self, msg='Test Message', toUserName=None): r = self.send_raw_msg(1, msg, toUserName) return r +def _prepare_file(fileDir, file_=None): + fileDict = {} + if file_: + if hasattr(file_, 'read'): + file_ = file_.read() + else: + return ReturnValue({'BaseResponse': { + 'ErrMsg': 'file_ param should be opened file', + 'Ret': -1005, }}) + else: + if not utils.check_file(fileDir): + return ReturnValue({'BaseResponse': { + 'ErrMsg': 'No file found in specific dir', + 'Ret': -1002, }}) + with open(fileDir, 'rb') as f: + file_ = f.read() + fileDict['fileSize'] = len(file_) + fileDict['fileMd5'] = hashlib.md5(file_).hexdigest() + fileDict['file_'] = io.BytesIO(file_) + return fileDict + def upload_file(self, fileDir, isPicture=False, isVideo=False, - toUserName='filehelper'): + toUserName='filehelper', file_=None, preparedFile=None): logger.debug('Request to upload a %s: %s' % ( 'picture' if isPicture else 'video' if isVideo else 'file', fileDir)) - if not utils.check_file(fileDir): - return ReturnValue({'BaseResponse': { - 'ErrMsg': 'No file found in specific dir', - 'Ret': -1002, }}) - fileSize = os.path.getsize(fileDir) + if not preparedFile: + preparedFile = _prepare_file(fileDir, file_) + if not preparedFile: + return preparedFile + fileSize, fileMd5, file_ = \ + preparedFile['fileSize'], preparedFile['fileMd5'], preparedFile['file_'] fileSymbol = 'pic' if isPicture else 'video' if isVideo else'doc' - with open(fileDir, 'rb') as f: - fileMd5 = hashlib.md5(f.read()).hexdigest() - file_ = open(fileDir, 'rb') chunks = int((fileSize - 1) / 524288) + 1 clientMediaId = int(time.time() * 1e4) uploadMediaRequest = json.dumps(OrderedDict([ @@ -313,7 +332,7 @@ def upload_file(self, fileDir, isPicture=False, isVideo=False, return ReturnValue(rawResponse=r) def upload_chunk_file(core, fileDir, fileSymbol, fileSize, - file, chunk, chunks, uploadMediaRequest): + file_, chunk, chunks, uploadMediaRequest): url = core.loginInfo.get('fileUrl', core.loginInfo['url']) + \ '/webwxuploadmedia?f=json' # save it on server @@ -331,7 +350,7 @@ def upload_chunk_file(core, fileDir, fileSymbol, fileSize, ('uploadmediarequest', (None, uploadMediaRequest)), ('webwx_data_ticket', (None, cookiesList['webwx_data_ticket'])), ('pass_ticket', (None, core.loginInfo['pass_ticket'])), - ('filename' , (os.path.basename(fileDir), file.read(524288), 'application/octet-stream'))]) + ('filename' , (os.path.basename(fileDir), file_.read(524288), 'application/octet-stream'))]) if chunks == 1: del files['chunk']; del files['chunks'] else: @@ -339,12 +358,21 @@ def upload_chunk_file(core, fileDir, fileSymbol, fileSize, headers = { 'User-Agent' : config.USER_AGENT } return requests.post(url, files=files, headers=headers) -def send_file(self, fileDir, toUserName=None, mediaId=None): +def send_file(self, fileDir, toUserName=None, mediaId=None, file_=None): logger.debug('Request to send a file(mediaId: %s) to %s: %s' % ( mediaId, toUserName, fileDir)) - if toUserName is None: toUserName = self.storageClass.userName + if hasattr(fileDir, 'read'): + return ReturnValue({'BaseResponse': { + 'ErrMsg': 'fileDir param should not be an opened file in send_file', + 'Ret': -1005, }}) + if toUserName is None: + toUserName = self.storageClass.userName if mediaId is None: - r = self.upload_file(fileDir) + preparedFile = _prepare_file(fileDir, file_) + if not preparedFile: + return preparedFile + fileSize = preparedFile['fileSize'] + r = self.upload_file(fileDir, preparedFile=preparedFile) if r: mediaId = r['MediaId'] else: @@ -354,10 +382,10 @@ def send_file(self, fileDir, toUserName=None, mediaId=None): 'BaseRequest': self.loginInfo['BaseRequest'], 'Msg': { 'Type': 6, - 'Content': ("%s"%os.path.basename(fileDir) + + 'Content': ("%s" % os.path.basename(fileDir) + "6" + - "%s%s"%(str(os.path.getsize(fileDir)), mediaId) + - "%s"%os.path.splitext(fileDir)[1].replace('.','')), + "%s%s" % (str(fileSize), mediaId) + + "%s" % os.path.splitext(fileDir)[1].replace('.','')), 'FromUserName': self.storageClass.userName, 'ToUserName': toUserName, 'LocalID': int(time.time() * 1e4), @@ -370,12 +398,22 @@ def send_file(self, fileDir, toUserName=None, mediaId=None): data=json.dumps(data, ensure_ascii=False).encode('utf8')) return ReturnValue(rawResponse=r) -def send_image(self, fileDir, toUserName=None, mediaId=None): +def send_image(self, fileDir=None, toUserName=None, mediaId=None, file_=None): logger.debug('Request to send a image(mediaId: %s) to %s: %s' % ( mediaId, toUserName, fileDir)) - if toUserName is None: toUserName = self.storageClass.userName + if fileDir or file_: + if hasattr(fileDir, 'read'): + file_, fileDir = fileDir, None + if fileDir is None: + fileDir = 'tmp.jpg' # specific fileDir to send gifs + else: + return ReturnValue({'BaseResponse': { + 'ErrMsg': 'Either fileDir or file_ should be specific', + 'Ret': -1005, }}) + if toUserName is None: + toUserName = self.storageClass.userName if mediaId is None: - r = self.upload_file(fileDir, isPicture=not fileDir[-4:] == '.gif') + r = self.upload_file(fileDir, isPicture=not fileDir[-4:] == '.gif', file_=file_) if r: mediaId = r['MediaId'] else: @@ -402,12 +440,22 @@ def send_image(self, fileDir, toUserName=None, mediaId=None): data=json.dumps(data, ensure_ascii=False).encode('utf8')) return ReturnValue(rawResponse=r) -def send_video(self, fileDir=None, toUserName=None, mediaId=None): +def send_video(self, fileDir=None, toUserName=None, mediaId=None, file_=None): logger.debug('Request to send a video(mediaId: %s) to %s: %s' % ( mediaId, toUserName, fileDir)) - if toUserName is None: toUserName = self.storageClass.userName + if fileDir or file_: + if hasattr(fileDir, 'read'): + file_, fileDir = fileDir, None + if fileDir is None: + fileDir = 'tmp.mp4' # specific fileDir to send other formats + else: + return ReturnValue({'BaseResponse': { + 'ErrMsg': 'Either fileDir or file_ should be specific', + 'Ret': -1005, }}) + if toUserName is None: + toUserName = self.storageClass.userName if mediaId is None: - r = self.upload_file(fileDir, isVideo=True) + r = self.upload_file(fileDir, isVideo=True, file_=file_) if r: mediaId = r['MediaId'] else: diff --git a/itchat/components/register.py b/itchat/components/register.py index 03add000..86a37fcc 100644 --- a/itchat/components/register.py +++ b/itchat/components/register.py @@ -66,7 +66,7 @@ def configured_reply(self): def msg_register(self, msgType, isFriendChat=False, isGroupChat=False, isMpChat=False): ''' a decorator constructor return a specific decorator based on information given ''' - if not isinstance(msgType, list): + if not (isinstance(msgType, list) or isinstance(msgType, tuple)): msgType = [msgType] def _msg_register(fn): for _msgType in msgType: @@ -78,6 +78,7 @@ def _msg_register(fn): self.functionDict['MpChat'][_msgType] = fn if not any((isFriendChat, isGroupChat, isMpChat)): self.functionDict['FriendChat'][_msgType] = fn + return fn return _msg_register def run(self, debug=False, blockThread=True): diff --git a/itchat/core.py b/itchat/core.py index 95ba1d52..ba384ced 100644 --- a/itchat/core.py +++ b/itchat/core.py @@ -313,7 +313,7 @@ def send_msg(self, msg='Test Message', toUserName=None): ''' raise NotImplementedError() def upload_file(self, fileDir, isPicture=False, isVideo=False, - toUserName='filehelper'): + toUserName='filehelper', file_=None, preparedFile=None): ''' upload file to server and get mediaId for options - fileDir: dir for file ready for upload @@ -325,7 +325,7 @@ def upload_file(self, fileDir, isPicture=False, isVideo=False, it is defined in components/messages.py ''' raise NotImplementedError() - def send_file(self, fileDir, toUserName=None, mediaId=None): + def send_file(self, fileDir, toUserName=None, mediaId=None, file_=None): ''' send attachment for options - fileDir: dir for file ready for upload @@ -335,7 +335,7 @@ def send_file(self, fileDir, toUserName=None, mediaId=None): it is defined in components/messages.py ''' raise NotImplementedError() - def send_image(self, fileDir, toUserName=None, mediaId=None): + def send_image(self, fileDir=None, toUserName=None, mediaId=None, file_=None): ''' send image for options - fileDir: dir for file ready for upload @@ -346,7 +346,7 @@ def send_image(self, fileDir, toUserName=None, mediaId=None): it is defined in components/messages.py ''' raise NotImplementedError() - def send_video(self, fileDir=None, toUserName=None, mediaId=None): + def send_video(self, fileDir=None, toUserName=None, mediaId=None, file_=None): ''' send video for options - fileDir: dir for file ready for upload