better error handle in job.py

This commit is contained in:
rustdesk 2024-03-16 00:30:30 +08:00
parent fbfcf862cb
commit a3f79065c5

View File

@ -37,7 +37,7 @@ def create(task_name, file_path=None):
headers=HEADERS, headers=HEADERS,
files=files, files=files,
) )
return response.json() return get_json(response)
def upload_file(task_id, file_path): def upload_file(task_id, file_path):
@ -46,12 +46,12 @@ def upload_file(task_id, file_path):
response = requests.post( response = requests.post(
f"{BASE_URL}/tasks/{task_id}/files", headers=HEADERS, files=files f"{BASE_URL}/tasks/{task_id}/files", headers=HEADERS, files=files
) )
return response.json() return get_json(response)
def get_status(task_id): def get_status(task_id):
response = requests.get(f"{BASE_URL}/tasks/{task_id}/status", headers=HEADERS) response = requests.get(f"{BASE_URL}/tasks/{task_id}/status", headers=HEADERS)
return response.json() return get_json(response)
def download_files(task_id, output_dir, fn=None): def download_files(task_id, output_dir, fn=None):
@ -88,7 +88,7 @@ def download_one_file(task_id, file_id, output_dir):
def fetch(): def fetch():
response = requests.get(f"{BASE_URL}/tasks/fetch_task", headers=HEADERS) response = requests.get(f"{BASE_URL}/tasks/fetch_task", headers=HEADERS)
return response.json() return get_json(response)
def update_status(task_id, status): def update_status(task_id, status):
@ -97,7 +97,7 @@ def update_status(task_id, status):
headers=HEADERS, headers=HEADERS,
json=status, json=status,
) )
return response.json() return get_json(response)
def delete_task(task_id): def delete_task(task_id):
@ -105,7 +105,7 @@ def delete_task(task_id):
f"{BASE_URL}/tasks/{task_id}", f"{BASE_URL}/tasks/{task_id}",
headers=HEADERS, headers=HEADERS,
) )
return response.json() return get_json(response)
def sign(file_path): def sign(file_path):
@ -152,6 +152,13 @@ def sign_one_file(file_path):
return False return False
def get_json(response):
try:
return response.json()
except Exception as e:
raise Exception(response.text)
SIGN_EXTENSIONS = [ SIGN_EXTENSIONS = [
".dll", ".dll",
".exe", ".exe",
@ -170,11 +177,18 @@ SIGN_EXTENSIONS = [
] ]
def sign_files(dir_path): def sign_files(dir_path, only_ext=None):
if only_ext:
only_ext = only_ext.split(",")
for i in range(len(only_ext)):
if not only_ext[i].startswith("."):
only_ext[i] = "." + only_ext[i]
for root, dirs, files in os.walk(dir_path): for root, dirs, files in os.walk(dir_path):
for file in files: for file in files:
file_path = os.path.join(root, file) file_path = os.path.join(root, file)
_, ext = os.path.splitext(file_path) _, ext = os.path.splitext(file_path)
if only_ext and ext not in only_ext:
continue
if ext in SIGN_EXTENSIONS: if ext in SIGN_EXTENSIONS:
if not sign_one_file(file_path): if not sign_one_file(file_path):
logging.error(f"Failed to sign {file_path}") logging.error(f"Failed to sign {file_path}")
@ -200,6 +214,9 @@ def main():
sign_files_parser.add_argument( sign_files_parser.add_argument(
"dir_path", help="The path of the directory containing the files to sign." "dir_path", help="The path of the directory containing the files to sign."
) )
sign_files_parser.add_argument(
"only_ext", help="The file extension to sign.", default=None, nargs="?"
)
# Create a parser for the "fetch" command # Create a parser for the "fetch" command
fetch_parser = subparsers.add_parser("fetch", help="Fetch a task.") fetch_parser = subparsers.add_parser("fetch", help="Fetch a task.")
@ -258,7 +275,7 @@ def main():
if args.command == "sign_one_file": if args.command == "sign_one_file":
sign_one_file(args.file_path) sign_one_file(args.file_path)
elif args.command == "sign_files": elif args.command == "sign_files":
sign_files(args.dir_path) sign_files(args.dir_path, args.only_ext)
elif args.command == "fetch": elif args.command == "fetch":
print(fetch()) print(fetch())
elif args.command == "update_status": elif args.command == "update_status":