api.py (6021B)
1 #!/usr/bin/python3 2 3 import asyncio 4 import collections 5 import enum 6 import io 7 import itertools 8 import pathlib 9 from typing import Any, AsyncIterator, Generic, Iterator, Optional, TypeVar 10 11 # normally I'd prefer using standard modules but streaming POST data is really important 12 import httpx 13 import msgspec 14 15 T = TypeVar("T") 16 17 GofileUpload = collections.namedtuple("GofileUpload", ["file", "result"]) 18 19 20 class GofileStatus(enum.Enum): 21 OK = "ok" 22 23 # header is incorrect 24 # this may happen if uploadFile was called without using multipart/form-data 25 ERROR_HEADERS = "headersError" 26 27 # no file was provided 28 # POST data may be malformed 29 ERROR_NOFILE = "error-noFile" 30 31 # no owner token was provided 32 # happens if a folderId was provided without the correct owner 33 ERROR_OWNER = "error-owner" 34 35 # token is invalid 36 # either a malformed token was initially provided or the guest token expired 37 ERROR_TOKEN = "error-token" 38 39 # token is not associated with a premium user 40 # as the API notes, guest tokens cannot access the API methods marked as premium 41 ERROR_NOT_PREMIUM = "error-notPremium" 42 43 # no server is available to process this request 44 ERROR_NO_SERVER = "noServer" 45 46 47 class GofileServerResult(msgspec.Struct): 48 server: Optional[str] = None 49 50 51 class GofileZonedServerResult(msgspec.Struct): 52 name: Optional[str] = None 53 zone: Optional[str] = None 54 55 def to_base_server(self): 56 return GofileServerResult(self.name) 57 58 59 class GofileServerListResult(msgspec.Struct): 60 servers: list[GofileZonedServerResult] 61 62 63 class GofileUploadResult(msgspec.Struct): 64 download_page: str = msgspec.field(name="downloadPage") 65 code: str = msgspec.field(name="parentFolderCode") 66 parent_folder: str = msgspec.field(name="parentFolder") 67 file_id: str = msgspec.field(name="id") 68 file_name: str = msgspec.field(name="name") 69 md5_hash: str = msgspec.field(name="md5") 70 servers: list[str] 71 72 # a guestToken field is provided if no access token was given and no folderID was specified 73 guest_token: Optional[str] = msgspec.field(default=None, name="guestToken") 74 75 76 class GofileServerResponse(msgspec.Struct, Generic[T]): 77 status: GofileStatus 78 data: T 79 80 81 async def _gofile_api_get(*args, type: type[T], **kwargs) -> T: 82 # performs a GET request and extracts the 'data' property from the response as a given type 83 # if the status is not 'ok', an exception is raised 84 async with httpx.AsyncClient() as client: 85 r = await client.get(*args, **kwargs) 86 r.raise_for_status() 87 88 # suppress the valid-type error since mypy doesn't cannot use runtime-specialized types 89 # but msgspec needs to do so 90 # https://stackoverflow.com/a/59636248 91 result = msgspec.json.decode(r.text, type=GofileServerResponse[type]) # type: ignore[valid-type] 92 93 if result.status != GofileStatus.OK: 94 raise Exception(result) 95 return result.data 96 97 98 async def _gofile_api_post(*args, type: type[T], **kwargs) -> T: 99 # performs a POST request and extracts the 'data' property from the response as a given type 100 # if the status is not 'ok', an exception is raised 101 async with httpx.AsyncClient() as client: 102 r = await client.post(*args, **kwargs) 103 r.raise_for_status() 104 105 # see typing woes at _gofile_api_get 106 result = msgspec.json.decode(r.text, type=GofileServerResponse[type]) # type: ignore[valid-type] 107 108 if result.status != GofileStatus.OK: 109 raise Exception(result) 110 return result.data 111 112 113 async def upload_single( 114 file: io.FileIO, 115 token: Optional[str] = None, 116 folder_id: Optional[str] = None, 117 server: Optional[str] = None, 118 ) -> GofileUpload: 119 """ 120 Uploads a single file. 121 122 :param file: An open file handle. 123 :param token: Token used for uploading. If not specified, the returned 124 ``GofileUpload.result.guest_token`` should be used for subsequent uploads to 125 the same folder. 126 :param folder_id: Folder to upload to. If not specified, the returned 127 ``GofileUpload.result.parent_folder`` should be used for subsequent 128 uploads to the same folder. 129 :param server: No longer used. Previously specified a specific subdomain to upload to. 130 """ 131 # we return a GofileUpload instead of a GofileUploadResult so there's consistency between upload_single / upload_multiple 132 133 # automatically shorten long file names in small terminals (e.g. split panes) 134 for attempt in itertools.count(1): 135 try: 136 filepath = pathlib.Path(str(file.name)) 137 post_data: dict[str, Any] = { 138 "file": (filepath.name, file, "application/octet-stream"), 139 } 140 141 if token: 142 post_data["token"] = token 143 if folder_id: 144 post_data["folderId"] = folder_id 145 146 upload_result = await _gofile_api_post( 147 "https://upload.gofile.io/uploadfile", files=post_data, type=GofileUploadResult 148 ) 149 break 150 except httpx.HTTPError as e: 151 print(e) 152 await asyncio.sleep(min(300, 10 * (1.2**attempt))) 153 pass 154 155 return GofileUpload(file, upload_result) 156 157 158 async def upload_multiple( 159 files: Iterator[io.FileIO], 160 token: Optional[str] = None, 161 folder_id: Optional[str] = None, 162 server: Optional[str] = None, 163 ) -> AsyncIterator[GofileUpload]: 164 """ 165 Uploads multiple files to the same folder, returning an interator of results. 166 """ 167 first_file, *other_files = files 168 169 first_upload = await upload_single(first_file, token, folder_id, server) 170 171 if not token: 172 token = first_upload.result.guest_token 173 if not folder_id: 174 folder_id = first_upload.result.parent_folder 175 176 yield first_upload 177 178 uploads = [upload_single(file, token, folder_id, server) for file in other_files] 179 for upload in asyncio.as_completed(uploads): 180 yield await upload