add an example script for scheduled quota updates

devopsmakers · Apr 26, 2021 · 32db078 · 32db078
1 parent 1275328
commit 32db078
Show file tree

Hide file tree

Showing 5 changed files with 143 additions and 4 deletions.
diff --git a/docs/full-configuration.md b/docs/full-configuration.md
@@ -205,7 +205,7 @@ The configuration file contains the following sections:
   - `password_caching`, boolean. Verifying argon2id passwords has a high memory and computational cost, verifying bcrypt passwords has a high computational cost, by enabling, in memory, password caching you reduce these costs. Default: `true`
   - `update_mode`, integer. Defines how the database will be initialized/updated. 0 means automatically. 1 means manually using the initprovider sub-command.
   - `skip_natural_keys_validation`, boolean. If `true` you can use any UTF-8 character for natural keys as username, admin name, folder name. These keys are used in URIs for REST API and Web admin. If `false` only unreserved URI characters are allowed: ALPHA / DIGIT / "-" / "." / "_" / "~". Default: `false`.
-  - `delayed_quota_update`, integer. This configuration parameter defines the number of seconds to accumulate quota updates. If there are a lot of close uploads, accumulating quota updates can save you many queries to the data provider. If you want to track quotas, a scheduled quota update is recommended in any case, the stored quota size may be incorrect for several reasons, such as an unexpected shutdown, temporary provider failures, file copied outside of SFTPGo, and so on. 0 means immediate quota update.
+  - `delayed_quota_update`, integer. This configuration parameter defines the number of seconds to accumulate quota updates. If there are a lot of close uploads, accumulating quota updates can save you many queries to the data provider. If you want to track quotas, a scheduled quota update is recommended in any case, the stored quota may be incorrect for several reasons, such as an unexpected shutdown while uploading files, temporary provider failures, files copied outside of SFTPGo, and so on. You could use the [quotascan example](../examples/quotascan) as a starting point. 0 means immediate quota update.
 - **"httpd"**, the configuration for the HTTP server used to serve REST API and to expose the built-in web interface
   - `bindings`, list of structs. Each struct has the following fields:
     - `port`, integer. The port used for serving HTTP requests. Default: 8080.

diff --git a/examples/quotascan/README.md b/examples/quotascan/README.md
@@ -0,0 +1,21 @@
+# Update user quota
+
+The `scanuserquota` example script shows how to use the SFTPGo REST API to update the users' quota.
+
+The stored quota may be incorrect for several reasons, such as an unexpected shutdown while uploading files, temporary provider failures, files copied outside of SFTPGo, and so on.
+
+A quota scan updates the number of files and their total size for the specified user and the virtual folders, if any, included in his quota.
+
+If you want to track quotas, a scheduled quota scan is recommended. You could use this example as a starting point.
+
+The script is written in Python and has the following requirements:
+
+- python3 or python2
+- python [Requests](https://requests.readthedocs.io/en/master/) module
+
+The provided example tries to connect to an SFTPGo instance running on `127.0.0.1:8080` using the following credentials:
+
+- username: `admin`
+- password: `password`
+
+Please edit the script according to your needs.
diff --git a/examples/quotascan/scanuserquota b/examples/quotascan/scanuserquota
@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+from datetime import datetime
+import sys
+import time
+
+import pytz
+import requests
+
+try:
+	import urllib.parse as urlparse
+except ImportError:
+	import urlparse
+
+# change base_url to point to your SFTPGo installation
+base_url = "http://127.0.0.1:8080"
+# set to False if you want to skip TLS certificate validation
+verify_tls_cert = True
+# set the credentials for a valid admin here
+admin_user = "admin"
+admin_password = "password"
+
+
+# set your update conditions here
+def needQuotaUpdate(user):
+	if user["status"] == 0:  # inactive user
+		return False
+	if user["quota_size"] == 0 and user["quota_files"] == 0:  # no quota restrictions
+		return False
+	return True
+
+
+class UpdateQuota:
+
+	def __init__(self):
+		self.limit = 100
+		self.offset = 0
+		self.access_token = ""
+		self.access_token_expiration = None
+
+	def printLog(self, message):
+		print("{} - {}".format(datetime.now(), message))
+
+	def checkAccessToken(self):
+		if self.access_token != "" and self.access_token_expiration:
+			expire_diff = self.access_token_expiration - datetime.now(tz=pytz.UTC)
+			# we don't use total_seconds to be python 2 compatible
+			seconds_to_expire = expire_diff.days * 86400 + expire_diff.seconds
+			if seconds_to_expire > 180:
+				return
+
+		auth = requests.auth.HTTPBasicAuth(admin_user, admin_password)
+		r = requests.get(urlparse.urljoin(base_url, "api/v2/token"), auth=auth, verify=verify_tls_cert)
+		if r.status_code != 200:
+			self.printLog("error getting access token: {}".format(r.text))
+			sys.exit(1)
+		self.access_token = r.json()["access_token"]
+		self.access_token_expiration = pytz.timezone("UTC").localize(datetime.strptime(r.json()["expires_at"],
+																					"%Y-%m-%dT%H:%M:%SZ"))
+
+	def getAuthHeader(self):
+		self.checkAccessToken()
+		return {"Authorization": "Bearer " + self.access_token}
+
+	def waitForQuotaUpdate(self, username):
+		while True:
+			auth_header = self.getAuthHeader()
+			r = requests.get(urlparse.urljoin(base_url, "api/v2/quota-scans"), headers=auth_header, verify=verify_tls_cert)
+			if r.status_code != 200:
+				self.printLog("error getting quota scans while waiting for {}: {}".format(username, r.text))
+				sys.exit(1)
+
+			scanning = False
+			for scan in r.json():
+				if scan["username"] == username:
+					scanning = True
+			if not scanning:
+				break
+			self.printLog("waiting for the quota scan to complete for user {}".format(username))
+			time.sleep(2)
+
+		self.printLog("quota update for user {} finished".format(username))
+
+	def updateUserQuota(self, username):
+		self.printLog("starting quota update for user {}".format(username))
+		auth_header = self.getAuthHeader()
+		r = requests.post(urlparse.urljoin(base_url, "api/v2/quota-scans"), headers=auth_header,
+						json={"username":username}, verify=verify_tls_cert)
+		if r.status_code != 202:
+			self.printLog("error starting quota scan for user {}: {}".format(username, r.text))
+			sys.exit(1)
+		self.waitForQuotaUpdate(username)
+
+	def updateUsersQuota(self):
+		while True:
+			self.printLog("get users, limit {} offset {}".format(self.limit, self.offset))
+			auth_header = self.getAuthHeader()
+			payload = {"limit":self.limit, "offset":self.offset}
+			r = requests.get(urlparse.urljoin(base_url, "api/v2/users"), headers=auth_header, params=payload,
+						verify=verify_tls_cert)
+			if r.status_code != 200:
+				self.printLog("error getting users: {}".format(r.text))
+				sys.exit(1)
+			users = r.json()
+			for user in users:
+				if needQuotaUpdate(user):
+					self.updateUserQuota(user["username"])
+				else:
+					self.printLog("user {} does not need a quota update".format(user["username"]))
+
+			self.offset += len(users)
+			if len(users) < self.limit:
+				break
+
+
+if __name__ == '__main__':
+	q = UpdateQuota()
+	q.updateUsersQuota()
diff --git a/ftpd/ftpd_test.go b/ftpd/ftpd_test.go
@@ -562,7 +562,7 @@ func TestBasicFTPHandling(t *testing.T) {
 	assert.Eventually(t, func() bool { return len(common.Connections.GetStats()) == 0 }, 1*time.Second, 50*time.Millisecond)
 }
 
-func TestLoginInvalidCredetials(t *testing.T) {
+func TestLoginInvalidCredentials(t *testing.T) {
 	u := getTestUser()
 	user, _, err := httpdtest.AddUser(u, http.StatusCreated)
 	assert.NoError(t, err)

diff --git a/httpd/schema/openapi.yaml b/httpd/schema/openapi.yaml
@@ -12,7 +12,7 @@ tags:
 info:
   title: SFTPGo
   description: SFTPGo REST API
-  version: 2.0.4
+  version: 2.0.5
   contact:
     url: 'https://github.com/drakkan/sftpgo'
   license:
@@ -320,7 +320,7 @@ paths:
       tags:
         - quota
       summary: Start user quota scan
-      description: Starts a new quota scan for the given user. A quota scan update the number of files and their total size for the specified user
+      description: Starts a new quota scan for the given user. A quota scan updates the number of files and their total size for the specified user and the virtual folders, if any, included in his quota
       operationId: start_quota_scan
       requestBody:
         required: true