crawlab_feader/FworkSpider/untils/attachment.py

# -*- coding: utf-8 -*-
"""
Created on 2020-09-06
---------
@summary: 附件下载模块
---------
@author: Dzr
"""
import io
import os
import uuid

import requests
import tqdm
import urllib3

import feapder.utils.tools as tools
from feapder.utils.log import log as logger
from untils.aliyun import AliYunService
from untils.execptions import AttachmentNullError

urllib3.disable_warnings()

headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.82 Safari/537.36",
    "Accept": "*/*"
}


class AttachmentDownloader:

    def __init__(self):
        self.dir_name = "file"

    def create_dir(self):
        if not os.path.exists(self.dir_name):
            os.makedirs(self.dir_name, mode=0o777, exist_ok=True)

    def create_file(self, filename, filetype):
        self.create_dir()
        sign = tools.get_sha1("{}_{}".format(filename, uuid.uuid4()))
        file_name = "{}.{}".format(sign, filetype)
        return "{}/{}".format(self.dir_name, file_name)

    @staticmethod
    def clean_attachment(file_path):
        """
        删除文件

        :param str file_path: 文件路径
        """
        try:
            os.remove(file_path)
        except FileNotFoundError:
            pass

    def remove(self, file):
        self.clean_attachment(file)

    @staticmethod
    def get_mb(data):
        """
        获取数据的Mb

        :param int data: 准备计算大小的内容
        :return: float
        """
        _kb = float(data / 1024.0)
        return float(_kb / 1024.0)

    @staticmethod
    def getsize(data):
        """
        计算数据大小

        :param data: 待上传的内容。
        :type data: bytes，str或file-like object
        :return str
        """
        size = 0
        if isinstance(data, str):
            try:
                size = os.path.getsize(data)
            except FileNotFoundError:
                pass
        elif isinstance(data, bytes):
            size = len(data)
        else:
            pass

        _kb = float(size) / 1024
        result = "{:.1f} kb".format(_kb)
        if _kb >= 1024:
            _M = _kb / 1024
            if _M >= 1024:
                _G = _M / 1024
                result = "{:.1f} G".format(_G)
            else:
                result = "{:.1f} M".format(_M)
        return result

    def fetch_data(self, url, proxies=None, file=None, show_error_log=False, **kwargs):
        """
        数据下载

        :param str url: 下载地址
        :param file: 本地文件
        :param dict kwargs: requests请求参数
        :param dict proxies: 代理ip
        :param show_error_log: 展示错误堆栈信息日志
        """
        method = kwargs.pop("method", "get")

        request_kwargs = {}
        request_kwargs.setdefault("proxies", proxies)
        request_kwargs.setdefault("headers", kwargs.get("headers") or headers)
        request_kwargs.setdefault("data", kwargs.pop("data", None))
        request_kwargs.setdefault("cookies", kwargs.pop("cookies", None))
        request_kwargs.setdefault("timeout", kwargs.pop("timeout", 60))
        request_kwargs.setdefault("stream", kwargs.pop("stream", True))
        request_kwargs.setdefault("verify", kwargs.pop("verify", False))
        request_kwargs.setdefault("allow_redirects", kwargs.pop("allow_redirects", True))

        retries = 0
        while retries < 3:
            try:
                with requests.request(method, url, **request_kwargs) as req:
                    stream = io.BytesIO()
                    lh = {k.lower(): v for k, v in req.headers.items()}
                    '''内容长度'''
                    cl = lh.get("content-length") or len(req.content)
                    icl = int(cl)
                    content_length = self.get_mb(icl)
                    if content_length > 50:
                        '''丢弃超过50Mb内容长度的文件'''
                        return stream.getvalue()

                    if req.status_code != 200:
                        retries += 1
                        continue

                    iter_content = req.iter_content(chunk_size=1024 * 20)
                    with tqdm.tqdm(
                            total=icl,
                            unit="B",
                            initial=0,
                            unit_scale=True,
                            unit_divisor=1024,  # 1M=1024Kb,单位换算
                            ascii=True,
                            desc=file) as bar:
                        if file is not None:
                            with open(file, "wb") as f:
                                for chunk in iter_content:
                                    stream.write(chunk)
                                    size = f.write(chunk)
                                    bar.update(size)
                        else:
                            for chunk in iter_content:
                                size = stream.write(chunk)
                                bar.update(size)
                    return stream.getvalue()
            except requests.RequestException as why:
                retries += 1
                if show_error_log:
                    logger.exception(why)

        return b''

    def _push_oss_from_stream(self, filename, filetype, url, **kwargs):
        """
        将数据流推送oss

        :param str filename: 文件名称
        :param str filetype: 文件类型
        :param str url: 下载地址
        """
        stream = self.fetch_data(url, file=None, **kwargs)
        attachment = {
            "filename": "{}.{}".format(filename, filetype),
            "org_url": url
        }
        if len(stream) > 0:
            fid = tools.get_sha1(stream)
            try:
                attachment["ftype"] = filetype
                attachment["fid"] = "{}.{}".format(fid, filetype)
                attachment["size"] = self.getsize(stream)
                attachment["url"] = "oss"
                AliYunService().push_oss_from_stream(attachment["fid"], stream)
            except Exception:
                pass

        return attachment

    def _push_oss_from_local(self, filename, filetype, url, **kwargs):
        """
        将本地文件推送oss

        :param str filename: 文件名称
        :param str filetype: 文件类型
        :param str url: 下载地址
        """
        file = self.create_file(filename, filetype)
        stream = self.fetch_data(url, file=file, **kwargs)
        '''上传/下载,无论失败成功都需要返回文件基础信息'''
        attachment = {
            "filename": "{}.{}".format(filename, filetype),
            "org_url": url
        }
        if len(stream) > 0:
            content_hash = tools.get_sha1(stream)
            try:
                attachment["fid"] = "{}.{}".format(content_hash, filetype)
                attachment["size"] = self.getsize(file)
                attachment["ftype"] = filetype
                attachment["url"] = "oss"
                AliYunService().push_oss_from_local(attachment["fid"], file)
            except Exception as e:
                logger.error(
                    "[{}]上传失败,原因:{}".format(filename, e.__class__.__name__)
                )

        self.remove(file)  # 删除本地临时文件
        return attachment

    def _fetch_attachment(self, filename, filetype, download_url, mode, **kwargs):
        """
        下载附件

        :param str filename: 文件名称
        :param str filetype: 文件类型
        :param str download_url: 下载地址
        :param str mode: 附件上传模式 "local" or "stream"
        """
        file_kwargs = dict(
            filename=filename,
            filetype=filetype,
            url=download_url,
            **kwargs
        )
        if mode == "stream":
            res = self._push_oss_from_stream(**file_kwargs)
        else:
            res = self._push_oss_from_local(**file_kwargs)
        return res

    def fetch_attachment(
        self,
        file_name: str,
        file_type: str,
        download_url: str,
        proxies=None,
        mode="local",
        **kwargs
    ):
        if not file_name or not file_type or not download_url:
            raise AttachmentNullError

        file_kwargs = dict(proxies=proxies, **kwargs)
        return self._fetch_attachment(file_name, file_type, download_url, mode, **file_kwargs)