280 lines
8.2 KiB
Python
280 lines
8.2 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
import requests
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
import argparse
|
||
|
|
import logging
|
||
|
|
import shutil
|
||
|
|
import zipfile
|
||
|
|
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.INFO,
|
||
|
|
format="%(asctime)s [%(levelname)s] %(message)s [%(filename)s:%(lineno)d]",
|
||
|
|
handlers=[logging.StreamHandler()],
|
||
|
|
)
|
||
|
|
|
||
|
|
# The URL of your Flask server
|
||
|
|
BASE_URL = os.getenv("BASE_URL") or "http://localhost:5000"
|
||
|
|
|
||
|
|
# The secret key for API authentication
|
||
|
|
SECRET_KEY = os.getenv("SECRET_KEY") or "worldpeace2024"
|
||
|
|
|
||
|
|
# The headers for API requests
|
||
|
|
HEADERS = {"Authorization": f"Bearer {SECRET_KEY}"}
|
||
|
|
|
||
|
|
TIMEOUT = int(os.getenv("TIMEOUT") or '30')
|
||
|
|
|
||
|
|
|
||
|
|
def create(task_name, file_path=None):
|
||
|
|
if file_path is None:
|
||
|
|
response = requests.post(f"{BASE_URL}/tasks/{task_name}", headers=HEADERS)
|
||
|
|
else:
|
||
|
|
with open(file_path, "rb") as f:
|
||
|
|
files = {"file": f}
|
||
|
|
response = requests.post(
|
||
|
|
f"{BASE_URL}/tasks/{task_name}",
|
||
|
|
headers=HEADERS,
|
||
|
|
files=files,
|
||
|
|
)
|
||
|
|
return response.json()
|
||
|
|
|
||
|
|
|
||
|
|
def upload_file(task_id, file_path):
|
||
|
|
with open(file_path, "rb") as f:
|
||
|
|
files = {"file": f}
|
||
|
|
response = requests.post(
|
||
|
|
f"{BASE_URL}/tasks/{task_id}/files", headers=HEADERS, files=files
|
||
|
|
)
|
||
|
|
return response.json()
|
||
|
|
|
||
|
|
|
||
|
|
def get_status(task_id):
|
||
|
|
response = requests.get(f"{BASE_URL}/tasks/{task_id}/status", headers=HEADERS)
|
||
|
|
return response.json()
|
||
|
|
|
||
|
|
|
||
|
|
def download_files(task_id, output_dir, fn=None):
|
||
|
|
response = requests.get(
|
||
|
|
f"{BASE_URL}/tasks/{task_id}/files", headers=HEADERS, stream=True
|
||
|
|
)
|
||
|
|
|
||
|
|
# Check if the request was successful
|
||
|
|
if fn is None:
|
||
|
|
fn = f"task_{task_id}_files.zip"
|
||
|
|
if response.status_code == 200:
|
||
|
|
# Save the file to the output directory
|
||
|
|
with open(os.path.join(output_dir, fn), "wb") as f:
|
||
|
|
for chunk in response.iter_content(chunk_size=1024):
|
||
|
|
if chunk:
|
||
|
|
f.write(chunk)
|
||
|
|
return response.ok
|
||
|
|
|
||
|
|
|
||
|
|
def download_one_file(task_id, file_id, output_dir):
|
||
|
|
response = requests.get(
|
||
|
|
f"{BASE_URL}/tasks/{task_id}/files/{file_id}", headers=HEADERS, stream=True
|
||
|
|
)
|
||
|
|
|
||
|
|
# Check if the request was successful
|
||
|
|
if response.status_code == 200:
|
||
|
|
# Save the file to the output directory
|
||
|
|
with open(os.path.join(output_dir, file_id), "wb") as f:
|
||
|
|
for chunk in response.iter_content(chunk_size=1024):
|
||
|
|
if chunk:
|
||
|
|
f.write(chunk)
|
||
|
|
return response.ok
|
||
|
|
|
||
|
|
|
||
|
|
def fetch():
|
||
|
|
response = requests.get(f"{BASE_URL}/tasks/fetch_task", headers=HEADERS)
|
||
|
|
return response.json()
|
||
|
|
|
||
|
|
|
||
|
|
def update_status(task_id, status):
|
||
|
|
response = requests.patch(
|
||
|
|
f"{BASE_URL}/tasks/{task_id}/status",
|
||
|
|
headers=HEADERS,
|
||
|
|
json=status,
|
||
|
|
)
|
||
|
|
return response.json()
|
||
|
|
|
||
|
|
|
||
|
|
def delete_task(task_id):
|
||
|
|
response = requests.delete(
|
||
|
|
f"{BASE_URL}/tasks/{task_id}",
|
||
|
|
headers=HEADERS,
|
||
|
|
)
|
||
|
|
return response.json()
|
||
|
|
|
||
|
|
|
||
|
|
def sign(file_path):
|
||
|
|
res = create("sign", file_path)
|
||
|
|
if res.ok:
|
||
|
|
task_id = res.task_id
|
||
|
|
|
||
|
|
# Poll the status every second
|
||
|
|
while True:
|
||
|
|
status = get_status(task_id)
|
||
|
|
if status["status"] == "done":
|
||
|
|
# Download the files
|
||
|
|
download_files(task_id, "output")
|
||
|
|
|
||
|
|
# Delete the task
|
||
|
|
delete_task(task_id)
|
||
|
|
|
||
|
|
break
|
||
|
|
|
||
|
|
time.sleep(1)
|
||
|
|
|
||
|
|
|
||
|
|
def sign_one_file(file_path):
|
||
|
|
logging.info(f"Signing {file_path}")
|
||
|
|
res = create("sign", file_path)
|
||
|
|
logging.info(f"Uploaded {file_path}")
|
||
|
|
task_id = res["id"]
|
||
|
|
n = 0
|
||
|
|
while True:
|
||
|
|
if n >= TIMEOUT:
|
||
|
|
delete_task(task_id)
|
||
|
|
logging.error(f"Failed to sign {file_path}")
|
||
|
|
break
|
||
|
|
time.sleep(6)
|
||
|
|
n += 1
|
||
|
|
status = get_status(task_id)
|
||
|
|
if status and status["state"] == "done":
|
||
|
|
download_one_file(
|
||
|
|
task_id, os.path.basename(file_path), os.path.dirname(file_path)
|
||
|
|
)
|
||
|
|
delete_task(task_id)
|
||
|
|
logging.info(f"Signed {file_path}")
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
SIGN_EXTENSIONS = [
|
||
|
|
".dll",
|
||
|
|
".exe",
|
||
|
|
".sys",
|
||
|
|
".vxd",
|
||
|
|
".msix",
|
||
|
|
".msixbundle",
|
||
|
|
".appx",
|
||
|
|
".appxbundle",
|
||
|
|
".msi",
|
||
|
|
".msp",
|
||
|
|
".msm",
|
||
|
|
".cab",
|
||
|
|
".ps1",
|
||
|
|
".psm1",
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def sign_files(dir_path):
|
||
|
|
for root, dirs, files in os.walk(dir_path):
|
||
|
|
for file in files:
|
||
|
|
file_path = os.path.join(root, file)
|
||
|
|
_, ext = os.path.splitext(file_path)
|
||
|
|
if ext in SIGN_EXTENSIONS:
|
||
|
|
if not sign_one_file(file_path):
|
||
|
|
logging.error(f"Failed to sign {file_path}")
|
||
|
|
break
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Command line interface for task operations."
|
||
|
|
)
|
||
|
|
subparsers = parser.add_subparsers(dest="command")
|
||
|
|
|
||
|
|
# Create a parser for the "sign_one_file" command
|
||
|
|
sign_one_file_parser = subparsers.add_parser(
|
||
|
|
"sign_one_file", help="Sign a single file."
|
||
|
|
)
|
||
|
|
sign_one_file_parser.add_argument("file_path", help="The path of the file to sign.")
|
||
|
|
|
||
|
|
# Create a parser for the "sign_files" command
|
||
|
|
sign_files_parser = subparsers.add_parser(
|
||
|
|
"sign_files", help="Sign all files in a directory."
|
||
|
|
)
|
||
|
|
sign_files_parser.add_argument(
|
||
|
|
"dir_path", help="The path of the directory containing the files to sign."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create a parser for the "fetch" command
|
||
|
|
fetch_parser = subparsers.add_parser("fetch", help="Fetch a task.")
|
||
|
|
|
||
|
|
# Create a parser for the "update_status" command
|
||
|
|
update_status_parser = subparsers.add_parser(
|
||
|
|
"update_status", help="Update the status of a task."
|
||
|
|
)
|
||
|
|
update_status_parser.add_argument("task_id", help="The ID of the task to update.")
|
||
|
|
update_status_parser.add_argument("status", help="The new status of the task.")
|
||
|
|
|
||
|
|
# Create a parser for the "delete_task" command
|
||
|
|
delete_task_parser = subparsers.add_parser("delete_task", help="Delete a task.")
|
||
|
|
delete_task_parser.add_argument("task_id", help="The ID of the task to delete.")
|
||
|
|
|
||
|
|
# Create a parser for the "create" command
|
||
|
|
create_parser = subparsers.add_parser("create", help="Create a task.")
|
||
|
|
create_parser.add_argument("task_name", help="The name of the task to create.")
|
||
|
|
create_parser.add_argument(
|
||
|
|
"file_path",
|
||
|
|
help="The path of the file for the task.",
|
||
|
|
default=None,
|
||
|
|
nargs="?",
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create a parser for the "upload_file" command
|
||
|
|
upload_file_parser = subparsers.add_parser(
|
||
|
|
"upload_file", help="Upload a file to a task."
|
||
|
|
)
|
||
|
|
upload_file_parser.add_argument(
|
||
|
|
"task_id", help="The ID of the task to upload the file to."
|
||
|
|
)
|
||
|
|
upload_file_parser.add_argument("file_path", help="The path of the file to upload.")
|
||
|
|
|
||
|
|
# Create a parser for the "get_status" command
|
||
|
|
get_status_parser = subparsers.add_parser(
|
||
|
|
"get_status", help="Get the status of a task."
|
||
|
|
)
|
||
|
|
get_status_parser.add_argument(
|
||
|
|
"task_id", help="The ID of the task to get the status of."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create a parser for the "download_files" command
|
||
|
|
download_files_parser = subparsers.add_parser(
|
||
|
|
"download_files", help="Download files from a task."
|
||
|
|
)
|
||
|
|
download_files_parser.add_argument(
|
||
|
|
"task_id", help="The ID of the task to download files from."
|
||
|
|
)
|
||
|
|
download_files_parser.add_argument(
|
||
|
|
"output_dir", help="The directory to save the downloaded files to."
|
||
|
|
)
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
if args.command == "sign_one_file":
|
||
|
|
sign_one_file(args.file_path)
|
||
|
|
elif args.command == "sign_files":
|
||
|
|
sign_files(args.dir_path)
|
||
|
|
elif args.command == "fetch":
|
||
|
|
print(fetch())
|
||
|
|
elif args.command == "update_status":
|
||
|
|
print(update_status(args.task_id, args.status))
|
||
|
|
elif args.command == "delete_task":
|
||
|
|
print(delete_task(args.task_id))
|
||
|
|
elif args.command == "create":
|
||
|
|
print(create(args.task_name, args.file_path))
|
||
|
|
elif args.command == "upload_file":
|
||
|
|
print(upload_file(args.task_id, args.file_path))
|
||
|
|
elif args.command == "get_status":
|
||
|
|
print(get_status(args.task_id))
|
||
|
|
elif args.command == "download_files":
|
||
|
|
print(download_files(args.task_id, args.output_dir))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|