Add bucket support

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<a id="create-bucket"></a>
+
+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<a id="delete-bucket"></a>
+
+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<a id="get-bucket"></a>
+
+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
+```

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" },
 ]

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

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

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")
+    )

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)