Merge pull request #17 from waifuvault/add-restriction

Add support for Restriction API

README.md

@@ -15,7 +15,7 @@ pip install waifuvault
 
 ## Usage
 
-This API contains 5 interactions:
+This API contains 9 interactions:
 
 1. [Upload File](#upload-file)
 2. [Get file Info](#get-file-info)
@@ -25,6 +25,7 @@ This API contains 5 interactions:
 6. [Create Bucket](#create-bucket)
 7. [Delete Bucket](#delete-bucket)
 8. [Get Bucket](#get-bucket)
+9. [Get Restrictions](#get-restrictions)
 
 The package is namespaced to `waifuvault`, so to import it, simply:
 
@@ -46,6 +47,7 @@ To Upload a file, use the `upload_file` function. This function takes the follow
 | `password`        | `string`           | If set, then the uploaded file will be encrypted                | false          |                                  |
 | `oneTimeDownload` | `boolean`          | if supplied, the file will be deleted as soon as it is accessed | false          |                                  |
 
+> **NOTE:** Server restrictions are checked by the SDK client side *before* upload, and will throw a ValueError exception if they are violated
 
 Using a URL:
 
@@ -236,4 +238,17 @@ import waifuvault
 bucket = waifuvault.get_bucket("some-bucket-token")
 print(bucket.token)
 print(bucket.files)  # Array of file objects
+```
+
+### Get Restrictions<a id="get-restrictions"></a>
+
+To get the list of restrictions applied to the server, you use the `get_restrictions` functions.
+
+This will respond with an array of name, value entries describing the restrictions applied to the server.
+
+```python
+import waifuvault
+restricions = waifuvault.get_bucket("some-bucket-token")
+
+print(restrictions.Restrictions)  # Array of restriction objects
 ```

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "waifuvault"
-version = "1.4.4"
+version = "1.4.5"
 authors = [
     { name="Walker Aldridge", email="walker@waifuvault.moe" },
 ]

src/waifuvault/__init__.py

@@ -1,2 +1,3 @@
-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
+from .waifumodels import FileResponse, FileUpload, BucketResponse, Restriction, RestrictionResponse
+from .waifuvault import (upload_file, file_info, get_file, delete_file, file_update, create_bucket, get_bucket,
+                         delete_bucket, get_restrictions)

src/waifuvault/waifumodels.py

@@ -1,9 +1,12 @@
 # Models for waifuVault
 import io
+import mimetypes
+import os
+from datetime import datetime, timedelta
 
 
 class FileUpload:
-    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):
+    def __init__(self, target: str | io.BytesIO, 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
@@ -32,22 +35,84 @@ def build_parameters(self):
 
 
 class FileOptions:
-    def __init__(self, hide_filename: bool = False, one_time_download: bool = False, protected: bool = False):
-        self.hideFilename = hide_filename
-        self.oneTimeDownload = one_time_download
-        self.protected = protected
+    def __init__(self, hide_filename: bool = False, one_time_download: bool = False, protected: bool = False, dict_obj: {} = None):
+        if dict_obj is not None:
+            self.hideFilename = dict_obj.get("hideFilename")
+            self.oneTimeDownload = dict_obj.get("oneTimeDownload")
+            self.protected = dict_obj.get("protected")
+        else:
+            self.hideFilename = hide_filename
+            self.oneTimeDownload = one_time_download
+            self.protected = protected
 
 
 class FileResponse:
-    def __init__(self, token: str = None, url: str = None, retention_period: any = None, bucket: str = None, options: FileOptions = None):
-        self.token = token
-        self.url = url
-        self.retentionPeriod = retention_period
-        self.bucket = bucket
-        self.options = options
+    def __init__(self, token: str = None, url: str = None, retention_period: str | int = None, bucket: str = None, options: FileOptions = None, dict_obj: {} = None):
+        if dict_obj is not None:
+            self.token = dict_obj.get("token")
+            self.url = dict_obj.get("url")
+            self.retentionPeriod = dict_obj.get("retentionPeriod")
+            self.bucket = dict_obj.get("bucket")
+            self.options = FileOptions(dict_obj=dict_obj["options"])
+        else:
+            self.token = token
+            self.url = url
+            self.retentionPeriod = retention_period
+            self.bucket = bucket
+            self.options = options
 
 
 class BucketResponse:
-    def __init__(self, token: str = None, files: list[FileResponse] = None):
-        self.token = token
-        self.files = files
+    def __init__(self, token: str = None, files: list[FileResponse] = None, dict_obj: {} = None):
+        if dict_obj is not None:
+            self.files = []
+            self.token = dict_obj.get("token")
+            for file in dict_obj.get("files"):
+                self.files.append(FileResponse(dict_obj=file))
+        else:
+            self.token = token
+            self.files = files
+
+
+class Restriction:
+    def __init__(self, type: str = None, value: str | int | list[str] = None, dict_obj: {} = None):
+        if dict_obj is not None:
+            self.type = dict_obj.get("type")
+            self.value = dict_obj.get("value")
+        else:
+            self.type = type
+            self.value = value
+
+    def passes(self, file: FileUpload):
+        if file.is_url():
+            return
+        match self.type:
+            case "MAX_FILE_SIZE":
+                if file.is_buffer():
+                    if file.target.getbuffer().nbytes > self.value:
+                        raise ValueError(f'File size {file.target.getbuffer().nbytes} is larger than max allowed {self.value}')
+                elif os.path.getsize(file.target) > self.value:
+                    raise ValueError(f'File size {os.path.getsize(file.target)} is larger than max allowed {self.value}')
+                return
+            case "BANNED_MIME_TYPE":
+                if file.is_buffer():
+                    mime_type, encoding = mimetypes.guess_type(file.target_name)
+                else:
+                    mime_type, encoding = mimetypes.guess_type(file.target)
+                if mime_type in self.value.split(','):
+                    raise ValueError(f'File MIME type {mime_type} is not allowed for upload')
+                return
+            case _:
+                raise NotImplementedError(f'Restriction type {self.type} is not implemented')
+
+
+class RestrictionResponse:
+    def __init__(self, restrictions: list[Restriction] = None, rest_obj: [] = None):
+        if rest_obj is not None:
+            self.Restrictions = []
+            for rest in rest_obj:
+                self.Restrictions.append(Restriction(dict_obj=rest))
+            self.Expires = datetime.now() + timedelta(minutes=10)
+        else:
+            self.Restrictions = restrictions
+            self.Expires = datetime.now() + timedelta(minutes=10)

src/waifuvault/waifuvault.py

@@ -1,21 +1,31 @@
 __base_url__ = "https://waifuvault.moe/rest"
+__restrictions = None
 
 import json
 import os
+from datetime import datetime
 from io import BytesIO
 
 import requests
 from requests_toolbelt import MultipartEncoder
 
-from .waifumodels import FileResponse, FileUpload, FileOptions, BucketResponse
+from .waifumodels import FileResponse, FileUpload, BucketResponse, RestrictionResponse
+
+
+# Get Restrictions
+def get_restrictions():
+    url = f"{__base_url__}/resources/restrictions"
+    response = requests.get(url)
+    __check_error(response, False)
+    return RestrictionResponse(rest_obj=json.loads(response.text))
 
 
 # Create Bucket
 def create_bucket():
     url = f"{__base_url__}/bucket/create"
     response = requests.get(url)
     __check_error(response, False)
-    return __bucket_to_obj(json.loads(response.text))
+    return BucketResponse(dict_obj=json.loads(response.text))
 
 
 # Delete Bucket
@@ -32,12 +42,13 @@ def get_bucket(token: str):
     data = {"bucket_token": token}
     response = requests.post(url, json=data)
     __check_error(response, False)
-    return __bucket_to_obj(json.loads(response.text))
+    return BucketResponse(dict_obj=json.loads(response.text))
 
 
 # Upload File
 def upload_file(file_obj: FileUpload):
     url = __base_url__
+    __check_restrictions(file_obj)
     if file_obj.bucket_token:
         url += f"/{file_obj.bucket_token}"
     fields = {}
@@ -66,7 +77,7 @@ def upload_file(file_obj: FileUpload):
         data=multipart_data,
         headers=header_data)
     __check_error(response, False)
-    return __dict_to_obj(json.loads(response.text))
+    return FileResponse(dict_obj=json.loads(response.text))
 
 
 # Update File
@@ -85,7 +96,7 @@ def file_update(token: str, password: str = None, previous_password: str = None,
         data=fields
     )
     __check_error(response, False)
-    return __dict_to_obj(json.loads(response.text))
+    return FileResponse(dict_obj=json.loads(response.text))
 
 
 # Get File Info
@@ -96,7 +107,7 @@ def file_info(token: str, formatted: bool):
         params={'formatted': 'true' if formatted else 'false'}
     )
     __check_error(response, False)
-    return __dict_to_obj(json.loads(response.text))
+    return FileResponse(dict_obj=json.loads(response.text))
 
 
 # Delete File
@@ -137,24 +148,12 @@ def __check_error(response: requests.models.Response, is_download: bool):
     return
 
 
-def __dict_to_obj(dict_obj: any):
-    return FileResponse(
-        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):
-    actual_files = []
-    for file in bucket_obj.get("files"):
-        actual_files.append(__dict_to_obj(file))
-    return BucketResponse(
-        bucket_obj.get("token"),
-        actual_files
-    )
+# Check file restrictions
+def __check_restrictions(file_obj: FileUpload):
+    global __restrictions
+    if __restrictions is None:
+        __restrictions = get_restrictions()
+    if __restrictions is not None and __restrictions.Expires < datetime.now():
+        __restrictions = get_restrictions()
+    for restriction in __restrictions.Restrictions:
+        restriction.passes(file_obj)

tests/test_waifuvault.py

@@ -23,12 +23,15 @@ def __init__(self, ok, text, content=None, code=None):
                                   '{"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)
+restrictions_response = response_mock(True,
+                                      '[{"type": "MAX_FILE_SIZE","value": 536870912},{"type": "BANNED_MIME_TYPE","value": "application/x-msdownload,application/x-executable"}]')
 
 
 # URL Upload Tests
 def test_upload_url(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = ok_response_numeric)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
     upload_file = waifuvault.FileUpload("https://walker.moe/assets/sunflowers.png", expires="10m")
 
     # When
@@ -49,6 +52,7 @@ def test_upload_url(mocker):
 def test_upload_bucket(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = ok_response_numeric)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
     upload_file = waifuvault.FileUpload("https://walker.moe/assets/sunflowers.png", expires="10m", bucket_token="test-bucket")
 
     # When
@@ -70,6 +74,7 @@ def test_upload_bucket(mocker):
 def test_upload_url_error(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = bad_request)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
 
     # When
     upload_file = waifuvault.FileUpload("https://walker.moe/assets/sunflowers.png", expires="10m")
@@ -82,6 +87,7 @@ def test_upload_url_error(mocker):
 def test_upload_file(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = ok_response_numeric)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
     upload_file = waifuvault.FileUpload("tests/testfile.png", expires="10m")
 
     # When
@@ -98,6 +104,7 @@ def test_upload_file(mocker):
 def test_upload_file_error(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = bad_request)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
 
     # When
     upload_file = waifuvault.FileUpload("tests/testfile.png", expires="10m")
@@ -110,6 +117,7 @@ def test_upload_file_error(mocker):
 def test_upload_buffer(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = ok_response_numeric)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
     with open("tests/testfile.png", "rb") as fh:
         buf = io.BytesIO(fh.read())
     upload_file = waifuvault.FileUpload(buf,"testfile_buf.png",expires="10m")
@@ -128,6 +136,7 @@ def test_upload_buffer(mocker):
 def test_upload_buffer_error(mocker):
     # Given
     mock_put = mocker.patch('requests.put', return_value = bad_request)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
     with open("tests/testfile.png", "rb") as fh:
         buf = io.BytesIO(fh.read())
 
@@ -139,6 +148,21 @@ def test_upload_buffer_error(mocker):
         upload_res = waifuvault.upload_file(upload_file)
 
 
+def test_upload_restriction_error(mocker):
+    # Given
+    mock_put = mocker.patch('requests.put', return_value = bad_request)
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
+    with open("tests/testfile.png", "rb") as fh:
+        buf = io.BytesIO(fh.read())
+
+    # When
+    upload_file = waifuvault.FileUpload(buf, "testfile_buf.exe", expires="10m")
+
+    # Then
+    with pytest.raises(Exception, match=re.escape('File MIME type application/x-msdownload is not allowed for upload')):
+        upload_res = waifuvault.upload_file(upload_file)
+
+
 def test_file_info(mocker):
     # Given
     mock_get = mocker.patch('requests.get', return_value = ok_response_human)
@@ -279,6 +303,19 @@ def test_delete_bucket(mocker):
     assert (del_bucket is True), "Delete Bucket did not return true"
 
 
+def test_get_restrictions(mocker):
+    # Given
+    mock_get = mocker.patch('requests.get', return_value=restrictions_response)
+
+    # When
+    restrictions = waifuvault.get_restrictions()
+
+    # Then
+    mock_get.assert_called_once_with('https://waifuvault.moe/rest/resources/restrictions')
+    assert (isinstance(restrictions, waifuvault.RestrictionResponse)), "Get Restrictions did not return a retriction response instance"
+    assert (len(restrictions.Restrictions) == 2), "Get Restrictions wrong number of restrictions"
+
+
 def test_url_args():
     # Given
     file_down = waifuvault.FileUpload("https://waifuvault.moe/test", expires="1d", password="testpassword", hidefilename=True, oneTimeDownload=True)