From 909fa9b3ff0d81bba9b63dfdccbd10f68cc54eef Mon Sep 17 00:00:00 2001 From: Walker Aldridge Date: Mon, 19 Aug 2024 22:22:05 -0500 Subject: [PATCH] Add bucket support --- README.md | 55 ++++++++++++++++++++++++++++ pyproject.toml | 2 +- src/waifuvault/__init__.py | 4 +- src/waifuvault/waifumodels.py | 12 +++++- src/waifuvault/waifuvault.py | 46 ++++++++++++++++++++--- tests/test_waifuvault.py | 69 +++++++++++++++++++++++++++++++++-- 6 files changed, 175 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 9128f84..11b3c55 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,9 @@ This API contains 5 interactions: 3. [Update file Info](#update-file-info) 4. [Delete File](#delete-file) 5. [Get File](#get-file) +6. [Create Bucket](#create-bucket) +7. [Delete Bucket](#delete-bucket) +8. [Get Bucket](#get-bucket) The package is namespaced to `waifuvault`, so to import it, simply: @@ -37,6 +40,7 @@ To Upload a file, use the `upload_file` function. This function takes the follow |-------------------|--------------------|-----------------------------------------------------------------|----------------|----------------------------------| | `target` | `string or buffer` | The target to upload can be a buffer, URL or filename | true | URL or file path | | `target_name` | `string` | The filename of the target if it is a buffer | true if buffer | Filename with extension | +| `bucket_token` | 'string' | Token for a bucket to upload the file into | false | Create bucket gives token | | `expires` | `string` | A string containing a number and a unit (1d = 1day) | false | Valid units are `m`, `h` and `d` | | `hideFilename` | `boolean` | If true, then the uploaded filename won't appear in the URL | false | Defaults to `false` | | `password` | `string` | If set, then the uploaded file will be encrypted | false | | @@ -182,3 +186,54 @@ file_enc_down = waifuvault.get_file(upload_enc_res,"your_password") print(file_enc_down.__sizeof__()) ``` +### Create Bucket + +Buckets are virtual collections that are linked to your IP and a token. When you create a bucket, you will receive a bucket token that you can use in Get Bucket to get all the files in that bucket + +> **NOTE:** Only one bucket is allowed per client IP address, if you call it more than once, it will return the same bucket token + +To create a bucket, use the `create_bucket` function. This function does not take any arguments. + +```python +import waifuvault +bucket = waifuvault.create_bucket() +print(bucket.token) +``` + +### Delete Bucket + +Deleting a bucket will delete the bucket and all the files it contains. + +> **IMPORTANT:** All contained files will be **DELETED** along with the Bucket! + +To delete a bucket, you must call the `deleteBucket` function with the following options as parameters: + +| Option | Type | Description | Required | Extra info | +|-------------|-----------|-----------------------------------|----------|-------------------| +| `token` | `string` | The token of the bucket to delete | true | | + +> **NOTE:** `deleteBucket` will only ever either return `true` or throw an exception if the token is invalid + +```python +import waifuvault +resp = waifuvault.delete_bucket("some-bucket-token") +print(resp) +``` + +### Get Bucket + +To get the list of files contained in a bucket, you use the `get_bucket` functions and supply the token. +This function takes the following options as parameters: + +| Option | Type | Description | Required | Extra info | +|-------------|-----------|-------------------------|----------|-------------------| +| `token` | `string` | The token of the bucket | true | | + +This will respond with the bucket and all the files the bucket contains. + +```python +import waifuvault +bucket = waifuvault.get_bucket("some-bucket-token") +print(bucket.token) +print(bucket.files) # Array of file objects +``` \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0caf268..0d4809f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "waifuvault" -version = "1.4.1" +version = "1.4.2" authors = [ { name="Walker Aldridge", email="walker@waifuvault.moe" }, ] diff --git a/src/waifuvault/__init__.py b/src/waifuvault/__init__.py index c428936..c47de93 100644 --- a/src/waifuvault/__init__.py +++ b/src/waifuvault/__init__.py @@ -1,2 +1,2 @@ -from .waifumodels import FileResponse, FileUpload -from .waifuvault import upload_file, file_info, get_file, delete_file, file_update +from .waifumodels import FileResponse, FileUpload, BucketResponse +from .waifuvault import upload_file, file_info, get_file, delete_file, file_update, create_bucket, get_bucket, delete_bucket diff --git a/src/waifuvault/waifumodels.py b/src/waifuvault/waifumodels.py index 04c4ca7..84e5e9b 100644 --- a/src/waifuvault/waifumodels.py +++ b/src/waifuvault/waifumodels.py @@ -3,9 +3,10 @@ class FileUpload: - def __init__(self, target: any, target_name: str = "unknown", expires: str = None, password: str = None, hidefilename: bool = False, oneTimeDownload: bool = False): + def __init__(self, target: any, target_name: str = "unknown", bucket_token: str = None, expires: str = None, password: str = None, hidefilename: bool = False, oneTimeDownload: bool = False): self.target = target self.target_name = target_name + self.bucket_token = bucket_token self.hidefilename = hidefilename self.one_time_download = oneTimeDownload self.expires = expires @@ -31,10 +32,11 @@ def build_parameters(self): class FileResponse: - def __init__(self, token: str = None, url: str = None, retention_period: any = None, options: any = None): + def __init__(self, token: str = None, url: str = None, retention_period: any = None, bucket: str = None, options: any = None): self.token = token self.url = url self.retentionPeriod = retention_period + self.bucket = bucket self.options = options @@ -43,3 +45,9 @@ def __init__(self, hide_filename: bool = False, one_time_download: bool = False, self.hideFilename = hide_filename self.oneTimeDownload = one_time_download self.protected = protected + + +class BucketResponse: + def __init__(self, token: str = None, files: any = None): + self.token = token + self.files = files diff --git a/src/waifuvault/waifuvault.py b/src/waifuvault/waifuvault.py index e59ccc6..2c93578 100644 --- a/src/waifuvault/waifuvault.py +++ b/src/waifuvault/waifuvault.py @@ -7,11 +7,39 @@ import requests from requests_toolbelt import MultipartEncoder -from .waifumodels import FileResponse, FileUpload, FileOptions +from .waifumodels import FileResponse, FileUpload, FileOptions, BucketResponse + + +# Create Bucket +def create_bucket(): + url = f"{__base_url__}/bucket/createBucket" + response = requests.get(url) + __check_error(response, False) + return __bucket_to_obj(json.loads(response.text)) + + +# Delete Bucket +def delete_bucket(token: str): + url = f"{__base_url__}/bucket/{token}" + response = requests.delete(url) + __check_error(response, False) + return True if response.text == "true" else False + + +# Get Bucket +def get_bucket(token: str): + url = f"{__base_url__}/bucket" + data = {"bucket_token": token} + response = requests.post(url, json=data) + __check_error(response, False) + return __bucket_to_obj(json.loads(response.text)) # Upload File def upload_file(file_obj: FileUpload): + url = __base_url__ + if file_obj.bucket_token: + url += f"/{file_obj.bucket_token}" fields = {} if file_obj.password: fields['password'] = file_obj.password @@ -33,7 +61,7 @@ def upload_file(file_obj: FileUpload): header_data = {'Content-Type': multipart_data.content_type} response = requests.put( - __base_url__, + url, params=file_obj.build_parameters(), data=multipart_data, headers=header_data) @@ -111,11 +139,19 @@ def __check_error(response: requests.models.Response, is_download: bool): def __dict_to_obj(dict_obj: any): return FileResponse( - dict_obj["token"], - dict_obj["url"], - dict_obj["retentionPeriod"], + dict_obj.get("token"), + dict_obj.get("url"), + dict_obj.get("retentionPeriod"), + dict_obj.get("bucket"), FileOptions( dict_obj["options"]["hideFilename"], dict_obj["options"]["oneTimeDownload"], dict_obj["options"]["protected"] )) + + +def __bucket_to_obj(bucket_obj: any): + return BucketResponse( + bucket_obj.get("token"), + bucket_obj.get("files") + ) diff --git a/tests/test_waifuvault.py b/tests/test_waifuvault.py index 1f6dc07..e33f4b3 100644 --- a/tests/test_waifuvault.py +++ b/tests/test_waifuvault.py @@ -16,11 +16,11 @@ def __init__(self, ok, text, content=None, code=None): # Mocked responses ok_response_numeric = response_mock(True, - '{"url":"https://waifuvault.moe/f/something", "token":"test-token", "retentionPeriod":100, "options":{"protected": false, "oneTimeDownload": false, "hideFilename": false}}') + '{"url":"https://waifuvault.moe/f/something", "token":"test-token", "bucket":"test-bucket", "retentionPeriod":100, "options":{"protected": false, "oneTimeDownload": false, "hideFilename": false}}') ok_response_numeric_protected = response_mock(True, - '{"url":"https://waifuvault.moe/f/something", "token":"test-token", "retentionPeriod":100, "options":{"protected": true, "oneTimeDownload": false, "hideFilename": false}}') + '{"url":"https://waifuvault.moe/f/something", "token":"test-token", "bucket":"test-bucket", "retentionPeriod":100, "options":{"protected": true, "oneTimeDownload": false, "hideFilename": false}}') ok_response_human = response_mock(True, - '{"url":"https://waifuvault.moe/f/something", "token":"test-token", "retentionPeriod":"10 minutes", "options":{"protected": false, "oneTimeDownload": false, "hideFilename": false}}') + '{"url":"https://waifuvault.moe/f/something", "token":"test-token", "bucket":"test-bucket", "retentionPeriod":"10 minutes", "options":{"protected": false, "oneTimeDownload": false, "hideFilename": false}}') bad_request = response_mock(False, '{"name": "BAD_REQUEST", "message": "Error Test", "status": 400}',code=400) @@ -46,6 +46,27 @@ def test_upload_url(mocker): assert (upload_res.retentionPeriod == 100), "Retention does not match" +def test_upload_bucket(mocker): + # Given + mock_put = mocker.patch('requests.put', return_value = ok_response_numeric) + upload_file = waifuvault.FileUpload("https://walker.moe/assets/sunflowers.png", expires="10m", bucket_token="test-bucket") + + # When + upload_res = waifuvault.upload_file(upload_file) + mock_put.assert_called_once_with( + 'https://waifuvault.moe/rest/test-bucket', + params={'expires': '10m'}, + data={'url': 'https://walker.moe/assets/sunflowers.png'}, + headers=None) + + # Then + assert (upload_res.url == "https://waifuvault.moe/f/something"), "URL does not match" + assert (upload_res.token == "test-token"), "Token does not match" + assert (upload_res.bucket == "test-bucket"), "Bucket does not match" + assert (upload_res.options.protected is False), "Protected does not match" + assert (upload_res.retentionPeriod == 100), "Retention does not match" + + def test_upload_url_error(mocker): # Given mock_put = mocker.patch('requests.put', return_value = bad_request) @@ -214,6 +235,48 @@ def test_download_error(mocker): file_down = waifuvault.get_file(waifuvault.FileResponse(url="https://waifuvault.moe/f/something"), "dangerWaifu") +def test_create_bucket(mocker): + # Given + mock_create = mocker.patch('requests.get', + return_value=response_mock(True, + '{"token": "test-bucket", "files":[]}')) + + # When + bucket = waifuvault.create_bucket() + + # Then + mock_create.assert_called_once_with('https://waifuvault.moe/rest/bucket/createBucket') + assert (bucket.token == "test-bucket"), "Create Bucket did not return bucket" + + +def test_get_bucket(mocker): + # Given + mock_get = mocker.patch('requests.post', + return_value=response_mock(True, + '{"token": "test-bucket", "files":[]}')) + + # When + bucket = waifuvault.get_bucket("test-bucket") + + # Then + mock_get.assert_called_once_with('https://waifuvault.moe/rest/bucket', json={'bucket_token': 'test-bucket'}) + assert (bucket.token == "test-bucket"), "Get Bucket did not return bucket" + + +def test_delete_bucket(mocker): + # Given + mock_del = mocker.patch('requests.delete', + return_value=response_mock(True, + 'true')) + + # When + del_bucket = waifuvault.delete_bucket("test-bucket") + + # Then + mock_del.assert_called_once_with('https://waifuvault.moe/rest/bucket/test-bucket') + assert (del_bucket is True), "Delete Bucket did not return true" + + def test_url_args(): # Given file_down = waifuvault.FileUpload("https://waifuvault.moe/test", expires="1d", password="testpassword", hidefilename=True, oneTimeDownload=True)