import hashlib import os import re import traceback import uuid import warnings from urllib.parse import urlparse, unquote import requests import urllib3 from config.load import headers from utils.aliyun import AliYunService from utils.execptions import AttachmentNullError from utils.socks5 import Proxy urllib3.disable_warnings() def sha1(val): _sha1 = hashlib.sha1() if isinstance(val, bytes): _sha1.update(str(val).encode("utf-8")) elif isinstance(val, str): _sha1.update(val.encode("utf-8")) return _sha1.hexdigest() def remove(file_path: str): os.remove(file_path) def getsize(file_path: str): try: return os.path.getsize(file_path) except FileNotFoundError: return 0 def discern_file_format(text): file_types = { 'pdf', 'doc', 'docx', 'rar', 'zip', 'gzzb', 'jpg', 'png', 'swf' } for file_type in file_types: all_file_format = [file_type, file_type.upper()] for t in all_file_format: result = re.match(f'.*{t}$', text, re.S) if result is not None: return t else: return None def extract_file_type(text): if text is None: return None return discern_file_format(text) def extract_file_name_by_href(href: str, file_type: str): """从url中抽取文件名称""" # 中文标点符号:[\u3002\uff1b\uff0c\uff1a\u201c\u201d\uff08\uff09\u3001\uff1f\u300a\u300b] # 中文字符:[\u4e00 -\u9fa5] zh_char_pattern = '[\u3002\uff1b\uff0c\uff1a\u201c\u201d\uff08\uff09\u3001\uff1f\u300a\u300b\u4e00-\u9fa5]+' parser = urlparse(href) query = (parser.query or parser.path) result = re.search(f'.*\\.{file_type}', query, re.S) if result is not None: encode_str = unquote(result.group()) name = re.search(zh_char_pattern, encode_str) if name is not None: return unquote(name.group()) return None def extract_file_name(text): file_type = discern_file_format(text) if file_type is not None: repl = '.{}'.format(file_type) text = text.replace(repl, '') return text def verify_file_name(name): if extract_file_type(name) is None: raise ValueError class AttachmentDownloader(AliYunService): def __init__(self): super(AttachmentDownloader, self).__init__() self.dir_name = 'file' def _create_file(self, filename, filetype): os.makedirs(self.dir_name, mode=0o777, exist_ok=True) file = "{filename}.{filetype}".format( filename=sha1("{}_{}".format(filename, uuid.uuid4())), filetype=filetype ) return "{}/{}".format(self.dir_name, file) @staticmethod def _create_fid(file_stream: bytes): return sha1(file_stream) @staticmethod def _origin_filename(fid: str, filetype: str): return "{}.{}".format(fid, filetype) @staticmethod def _file_size(file: str): _kb = float(getsize(file)) / 1024 if _kb >= 1024: _M = _kb / 1024 if _M >= 1024: _G = _M / 1024 return "{:.1f} G".format(_G) else: return "{:.1f} M".format(_M) else: return "{:.1f} kb".format(_kb) @staticmethod def _download( url: str, file: str, enable_proxy=False, allow_show_exception=False, **kwargs ): request_params = {} request_params.setdefault('headers', kwargs.get('headers') or headers) request_params.setdefault('proxies', kwargs.get('proxies')) request_params.setdefault('timeout', kwargs.get('timeout') or 60) request_params.setdefault('stream', kwargs.get('stream') or True) request_params.setdefault('verify', kwargs.get('verify') or False) proxy = Proxy(enable_proxy) retries = 0 while retries < 3: try: with requests.get(url, **request_params) as req: if req.status_code == 200: stream = req.content with open(file, 'wb') as f: f.write(stream) return stream else: retries += 1 except requests.RequestException: if allow_show_exception: traceback.print_exc() if enable_proxy: proxy.switch() request_params.update({'proxies': proxy.proxies}) retries += 1 return b'' def download( self, file_name: str, file_type: str, download_url: str, enable_proxy=False, allow_request_exception=False, **kwargs ): if not file_name or not file_type or not download_url: raise AttachmentNullError local_tmp_file = self._create_file(file_name, file_type) file_stream = self._download( download_url, local_tmp_file, enable_proxy, allow_request_exception, **kwargs ) result = { 'filename': '{}.{}'.format(file_name, file_type), 'org_url': download_url } if len(file_stream) > 0: try: fid = self._create_fid(file_stream) key = self._origin_filename(fid, file_type) result.setdefault('fid', key) result.setdefault('ftype', file_type) result.setdefault('size', self._file_size(local_tmp_file)) result.setdefault('url', 'oss') super()._push_oss_from_local(key, local_tmp_file) except Exception as e: warnings.warn( "[{}]下载异常,原因:{}".format(file_name, e.__class__.__name__) ) remove(local_tmp_file) '''上传/下载,无论失败/成功必须返回附件信息''' return result