match_spider/zbytb/utils/attachment.py

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