70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
import argparse
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from rich.console import Console
|
|
from rich.progress import Progress
|
|
|
|
from .host_manager import host_manager
|
|
from .types import File, ProgressHandler
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description="Mirror a file on multiple file hosting services"
|
|
)
|
|
parser.add_argument(
|
|
"--pool",
|
|
type=int,
|
|
metavar="N",
|
|
default=5,
|
|
help="Number of max workers in the thread pool",
|
|
)
|
|
parser.add_argument("file", help="Path to the file to upload")
|
|
|
|
for host in host_manager.hosts:
|
|
parser.add_argument(
|
|
f"-{host.short}",
|
|
action="count",
|
|
default=0,
|
|
help=f"Upload to {host.site}. Repeat to upload multiple times",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
console = Console()
|
|
|
|
file_name = args.file
|
|
if not os.path.isfile(file_name):
|
|
console.log(f"[bold red]{repr(file_name)} is not a file!")
|
|
return 1
|
|
with open(file_name) as fhandle:
|
|
file_size = fhandle.seek(0, 2)
|
|
|
|
with Progress(console=console, transient=True) as progress, ThreadPoolExecutor(
|
|
max_workers=args.pool
|
|
) as thread_pool:
|
|
for host in host_manager.hosts:
|
|
for _ in range(getattr(args, host.short)):
|
|
thread_pool.submit(
|
|
host(
|
|
File(open(file_name, "rb"), file_name, file_size),
|
|
ProgressHandler(
|
|
console,
|
|
progress,
|
|
progress.add_task(
|
|
f"[{host.short}] Waiting...",
|
|
start=False,
|
|
total=file_size,
|
|
visible=False,
|
|
),
|
|
),
|
|
).upload
|
|
)
|
|
|
|
console.log("[bold green]File mirrored!")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|