SmartChat/file.py

import abc
import hashlib
import os
import random
from abc import ABC

import requests
import json


class BaseFileManager(ABC):

    @staticmethod
    def _post(url: str, body: dict, headers: dict = None):
        if headers is None:
            headers = {}
        resp = requests.post(url, json=body, headers=headers)
        if resp.status_code != 200:
            return {}

        return resp.json()

    @staticmethod
    def _put(url: str, raw: bytes, headers: dict = None):
        if headers is None:
            headers = {}
        resp = requests.put(url, data=raw, headers=headers)
        if resp.status_code not in [201, 200]:
            return {}
        return resp.json()

    @staticmethod
    def _get(url: str, headers: dict = None):
        if headers is None:
            headers = {}
        resp = requests.get(url, headers=headers)
        if resp.status_code != 200:
            return b""
        return resp.content

    @staticmethod
    def _md5(content: str):
        return hashlib.md5(content.encode("utf-8")).hexdigest()

    @abc.abstractmethod
    def upload_file(self, file_path: str, file_content: bytes):
        pass

    @abc.abstractmethod
    def get_file(self, file_path: str):
        pass


class ObjectFileManager(BaseFileManager):
    def __init__(self, endpoint: str = 'https://a.cmdp.cn/basiceg/v1/bytes'):
        self.endpoint = endpoint

    def create_ticket(self, file_path: str, file_name: str) -> str:
        url = "https://a.cmdp.cn/v1/cloudfile/file/createTicket"
        body = {
            "appId": "",
            "uri": f"s3://{file_path}",
            "startTime": "",
            "endTime": "",
            "maxTimes": 0,
            "uid": "",
            "name": file_name,
            "isPublic": False
        }
        response = self._post(url, body)
        return response.get("data", "")

    def upload_file(self, file_path: str, file_content: bytes):
        url = f"{self.endpoint}/{file_path}"
        response = self._put(url, file_content)
        return response

    def get_file(self, file_path: str):
        url = f"{self.endpoint}/{file_path}"
        response = self._get(url)
        return response

    @staticmethod
    def _get_file_list_info(file_path: str, page_index: int = 1, page_size: int = 1):
        url = "https://openapi.genesysinfo.net/gate/09/sys/key/list"

        payload = json.dumps({
            "path": f"s3://{file_path}",
            "pageIndex": page_index,
            "pageSize": page_size
        })
        headers = {
            'Content-Type': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload)
        if response.status_code != 200:
            return {}
        return response.json()

    def get_file_list(self, file_path: str):
        result = []
        file_list_info = self._get_file_list_info(file_path)
        total = file_list_info['totalRecords']
        for i in range(int(total / 500) + 1):
            result += self._get_file_list_info(file_path, page_index=i + 1, page_size=500)['modelList']
        return result

    def get_3party_file(self, url: str):
        response = self._get(url)
        return response

    def transfer_file(self, url: str, file_path: str, file_type: str):
        file = self.get_3party_file(url)
        file_name = self._md5(url)
        file_info = self.upload_file(f"{file_path}/{file_name}.{file_type}", file)
        temp_path = file_info.get("file", "")
        return temp_path

    def transfer_ticket(self, url: str, file_path: str, file_type: str):
        temp_path = self.transfer_file(url, file_path, file_type)
        if temp_path == "":
            return ""
        return self.create_ticket(temp_path, temp_path.split("/")[-1])


class LocalFileManager(BaseFileManager):
    def __init__(self, endpoint: str):
        os.makedirs(endpoint, exist_ok=True)
        self.endpoint = endpoint

    def _entire_path(self, path: str):
        return os.path.join(self.endpoint, path)

    @staticmethod
    def _exists(file_path: str):
        return os.path.exists(file_path)

    @staticmethod
    def _read(file_path: str):

        with open(file_path, "rb") as f:
            content = f.read()
        return content

    @staticmethod
    def _write(file_path: str, file_content: bytes):
        try:
            with open(file_path, "wb") as f:
                f.write(file_content)
            return True
        except Exception as e:
            print(e)
            return False

    def exists(self, file_path: str):
        return self._exists(self._entire_path(file_path))

    def get_size(self, file_path: str):
        file_path = self._entire_path(file_path)
        if not self._exists(file_path):
            return 0
        return os.path.getsize(file_path)

    def get_md5(self, file_path: str):
        file_path = self._entire_path(file_path)
        md5_hash = hashlib.md5()
        with open(file_path, "rb") as f:
            # 分块读取大文件，避免内存溢出
            for chunk in iter(lambda: f.read(4096), b""):
                md5_hash.update(chunk)
        md5_checksum = md5_hash.hexdigest()
        return md5_checksum

    def upload_file(self, file_path: str, file_content: bytes):
        entire_path = self._entire_path(file_path)
        return self._write(entire_path, file_content)

    def get_file(self, file_path: str):
        entire_path = self._entire_path(file_path)
        return self._read(entire_path)


class OptionChoice:
    def __init__(self, choices: list):
        self.choices = choices
        self.last_choices = []

    def _update_last_choices(self, choices):
        if len(self.last_choices) >= 39:
            self.last_choices = self.last_choices[3:]
        self.last_choices += choices

    def choose(self, num: int) -> list:
        temp = random.choices(self.choices, k=num)
        return temp