match_spider/zbytb/utils/attachment.py

import io
import os
import traceback
import uuid
from io import BytesIO

import oss2
import requests
import urllib3

from config.load import headers, ali_oss
from utils.clean_file import (
    clean_file_name,
    judge_file_url,
    modify_file_url_list,
    remove,
    limit_file_size
)
from utils.clean_file import sha1, getsize
from utils.log import logger
from utils.socks5 import Proxy

urllib3.disable_warnings()


class OssClient(object):
    def __init__(self, domain: str):
        # 初始化函数，用于创建类的实例
        self.domain = domain

    def upload(self, args: dict, retries=3, err_show=False) -> dict:
        reply = {"error_code": -1}

        raise_err = None
        for _ in range(retries):
            try:
                files = {
                    'file': (args['object_name'], BytesIO(args['stream'])),
                }
                data = {
                    'bucket_id': args['bucket_id'],
                    'object_name': args['object_name'],
                    'gzip': args.get('gzip', False),
                }

                response = requests.post(
                    f"{self.domain}/ossservice/upload",
                    files=files,
                    data=data,
                    timeout=300,
                )
                if response.status_code == 200:
                    reply.update(response.json())
                else:
                    reply['error_msg'] = f"HTTP error: {response.status_code}"

                raise_err = None
                break
            except Exception as e:
                reply['error_msg'] = str(e)
                raise_err = e

        if err_show and raise_err is not None:
            raise raise_err

        return reply

    def download(self, args: dict):
        reply = {}
        try:
            data = {
                "bucket_id": args["bucket_id"],
                "object_name": args["object_name"]
            }

            url = f"{self.domain}/ossservice/download"
            response = requests.post(url, data=data, timeout=300)
            response.raise_for_status()

            reply["error_code"] = 0
            reply["error_msg"] = "下载成功"
            reply["data"] = response.content

        except Exception as e:
            reply["error_code"] = -1
            reply["error_msg"] = str(e)

        return reply

    def delete(self, args: dict):
        reply = {}
        try:
            data = {
                "bucket_id": args["bucket_id"],
                "object_name": args["object_name"]
            }

            url = f"{self.domain}/ossservice/delete"
            response = requests.post(url, data=data, timeout=10)
            response.raise_for_status()

            reply = response.json()
            reply["error_code"] = 0
        except Exception as e:
            reply["error_code"] = -1
            reply["error_msg"] = str(e)
        return reply


class OssBucketClient:

    def __init__(self):
        key_id = ali_oss['key_id']
        key_secret = ali_oss['key_secret']
        endpoint = ali_oss['endpoint']
        bucket_name = ali_oss['bucket_name']
        auth = oss2.Auth(key_id, key_secret)
        self._bucket = oss2.Bucket(auth, endpoint, bucket_name)

    def push_oss_from_local(self, key, filename):
        """
        上传一个本地文件到OSS的普通文件

        :param str key: 上传到OSS的文件名
        :param str filename: 本地文件名，需要有可读权限
        """
        return self._bucket.put_object_from_file(key, filename)

    def push_oss_from_stream(self, key, data):
        """
        流式上传oss

        :param str key: 上传到OSS的文件名
        :param data: 待上传的内容。
        :type data: bytes，str或file-like object
        """
        return self._bucket.put_object(key, data)


class AttachmentDownloader:

    def __init__(self, address=None, mode=None):
        self.dir_name = 'file'
        if address is None:
            address = ali_oss['jy']['address']

        if mode == 'test':
            address = ali_oss['jy']['address_test']

        # self._oss = OssClient(domain=address)
        self._bucket = OssBucketClient()

    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 _calc_size(size: float):
        if size >= 1024:
            _M = size / 1024
            if _M >= 1024:
                _G = _M / 1024
                return "{:.1f} G".format(_G)
            else:
                return "{:.1f} M".format(_M)
        else:
            return "{:.1f} kb".format(size)

    def _file_size(self, file):
        _kb = float(getsize(file)) / 1024
        return self._calc_size(_kb)

    def stream_download(self, file_path, response):
        max_size = 1024 * 1024  # 1M
        stream = io.BytesIO()
        h = {k.lower(): v for k, v in response.headers.items()}
        content_length = h.get('content-length')
        if content_length is not None:
            content_size = int(content_length)
            content_length = float(content_size) / max_size  # 内容体总大小,单位：M
            if content_length > 50:
                # 丢弃超过50Mb内容长度的文件
                return stream.getvalue()
        else:
            content_size = None

        chunk_size = 1024  # 单次请求数据块最大值
        data_chunk_size = 0  # 下载数据块
        with open(file_path, 'wb') as f:
            for data in response.iter_content(chunk_size=chunk_size):
                size = stream.write(data)
                f.write(data)
                data_chunk_size += size
                if data_chunk_size / max_size > 50:
                    stream.truncate(0)  # 清空流
                    stream.seek(0)  # 将位置指针移回流的开始处
                    logger.warning("超过50M的附件|丢弃下载")
                    break

                # 如果content_size未设置，进行估算
                if content_size is None:
                    if size == chunk_size:
                        # 接收到完整块，假设总大小是当前的2倍
                        content_size = data_chunk_size * 2
                    else:
                        # 接收到最后一个块（通常小于chunk_size）
                        content_size = data_chunk_size

                # 计算并显示进度
                percent = min(100, int((data_chunk_size / content_size) * 100)) if content_size else 0
                print("\r文件下载进度：%d%%(%d/%d) - %s" % (percent, data_chunk_size, content_size, file_path), end=" ")

        return stream.getvalue()

    def _fetch_file(
        self,
        method: str,
        url: str,
        file: str,
        enable_proxy: bool,
        allow_show_exception: bool,
        **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.request(method, url, **request_params) as req:
                    req.raise_for_status()
                    return self.stream_download(file, req)
            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
    ):
        request_method = kwargs.pop('method', 'get')
        file_type = file_type.strip()
        file_name = clean_file_name(file_name, file_type)
        download_url = judge_file_url(download_url)

        for app_param in modify_file_url_list:
            download_url = app_param(download_url)

        local_tmp_file = self._create_file(file_name, file_type)
        file_stream = self._fetch_file(
            request_method,
            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')
                self._bucket.push_oss_from_local(key, local_tmp_file)

                # args = {
                #     "bucket_id": "file",
                #     "object_name": key,
                #     "gzip": False,
                #     "stream": file_stream
                # }
                # self._oss.upload(args, err_show=True)
            except Exception as e:
                logger.warning(
                    "[{}]下载异常,原因:{}".format(file_name, type(e).__name__)
                )

        remove(local_tmp_file)

        '''上传/下载,无论失败/成功必须返回附件信息'''
        if "size" not in result:
            return result
        elif limit_file_size(result.get('size')):
            return result
        else:
            return {}