mirror of
https://github.com/drakkan/sftpgo.git
synced 2024-11-22 15:40:23 +00:00
119 lines
3.7 KiB
Python
Executable file
119 lines
3.7 KiB
Python
Executable file
#!/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, timeout=10)
|
|
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/quotas/users/scans"), headers=auth_header, verify=verify_tls_cert,
|
|
timeout=10)
|
|
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/quotas/users/" + username + "/scan"), headers=auth_header,
|
|
verify=verify_tls_cert, timeout=10)
|
|
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, timeout=10)
|
|
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()
|