From aadb48c2b6fc18949bd485356bc2e7939bd9e0ed Mon Sep 17 00:00:00 2001 From: BarsTiger Date: Tue, 22 Feb 2022 18:20:34 +0200 Subject: [PATCH] New progress bars in CLI Now all changes are working Bins not updated --- modules/download.py | 76 +++++++++++++++++++++++++++++++++---------- modules/manager.py | 55 ++++++++++++++++++++----------- modules/virustotal.py | 3 +- 3 files changed, 96 insertions(+), 38 deletions(-) diff --git a/modules/download.py b/modules/download.py index a33784f..0eb435b 100644 --- a/modules/download.py +++ b/modules/download.py @@ -1,20 +1,62 @@ -import threading -from tqdm import tqdm -import requests +import os.path +from concurrent.futures import ThreadPoolExecutor +import signal +from functools import partial +from threading import Event +from urllib.request import urlopen + +from rich.progress import ( + BarColumn, + DownloadColumn, + Progress, + TaskID, + TextColumn, + TimeRemainingColumn, + TransferSpeedColumn, +) + +progress = Progress( + TextColumn("[bold blue]{task.fields[filename]}", justify="right"), + BarColumn(bar_width=None), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", + DownloadColumn(), + "•", + TransferSpeedColumn(), + "•", + TimeRemainingColumn(), +) -def dl(urls, save_to): - for url in urls: - def dl_thread(url_in_thread): - with requests.get(url_in_thread, stream=True) as r: - with open(save_to + "/" + url.split('/')[-1], "wb") as f: - pbar = tqdm(unit="B", unit_scale=True, total=int(r.headers['Content-Length']), - position=urls.index(url_in_thread)) - for chunk in r.iter_content(chunk_size=1024): - if chunk: - if pbar.n < pbar.total - 1: - pbar.update(len(chunk)) - f.write(chunk) - pbar.clear() +done_event = Event() - threading.Thread(target=dl_thread, args=(url,)).start() + +def handle_sigint(signum, frame): + done_event.set() + + +signal.signal(signal.SIGINT, handle_sigint) + + +def copy_url(task_id: TaskID, url: str, path: str) -> None: + progress.console.log(f"Requesting {url}") + response = urlopen(url) + progress.update(task_id, total=int(response.info()["Content-length"])) + with open(path, "wb") as dest_file: + progress.start_task(task_id) + for data in iter(partial(response.read, 32768), b""): + dest_file.write(data) + progress.update(task_id, advance=len(data)) + if done_event.is_set(): + return + progress.console.log(f"Downloaded {path}") + + +def dl(urls, dest_dir: str): + with progress: + with ThreadPoolExecutor(max_workers=len(urls)) as pool: + for url in urls: + filename = url.split("/")[-1] + dest_path = os.path.join(dest_dir, filename) + task_id = progress.add_task("download", filename=filename, start=False) + pool.submit(copy_url, task_id, url, dest_path) diff --git a/modules/manager.py b/modules/manager.py index 80f6390..26139e4 100644 --- a/modules/manager.py +++ b/modules/manager.py @@ -16,25 +16,34 @@ def install(package): r = r.text r = json.loads(r) - if r_code[1] not in [403, 401]: + if r_code[1] not in [200, 201]: return r_code[1] try: + # Inform the user print(f"[green]App {r['name']} found, information loaded[/]") + print() - if not os.path.exists('{1}apps/{0}'.format(r['name'], horsy_vars.horsypath)): - os.makedirs('{1}apps/{0}'.format(r['name'], horsy_vars.horsypath)) + # Create the app directory + if not os.path.exists('{1}apps\{0}'.format(r['name'], horsy_vars.horsypath)): + os.makedirs('{1}apps\{0}'.format(r['name'], horsy_vars.horsypath)) + # Get all download files urls print(f"Downloading {r['url'].split('/')[-1]}") to_download = [r['url']] if r['download']: print(f"Found dependency") - print(f"Downloading {r['download'].split('/')[-1]}") to_download.append(r['download']) - dl(to_download, '{0}apps/{1}'.format(horsy_vars.horsypath, r['name'])) - scan_to_cli('{2}apps/{0}/{1}'.format(r['name'], r['url'].split('/')[-1], horsy_vars.horsypath)) + # Download all files + dl(to_download, '{0}apps\{1}'.format(horsy_vars.horsypath, r['name'])) + print() + # Scan main file + scan_to_cli('{2}apps\{0}\{1}'.format(r['name'], r['url'].split('/')[-1], horsy_vars.horsypath)) + print() + + # Unzip the main file if needed def unzip(file, where): with zipfile.ZipFile(file, 'r') as zip_ref: zip_ref.extractall(where) @@ -42,25 +51,33 @@ def install(package): if r['url'].split('.')[-1] == 'zip': print(f"Extracting {r['url'].split('/')[-1]}") - unzip('{2}apps/{0}/{1}'.format(r['name'], r['url'].split('/')[-1], horsy_vars.horsypath), - '{1}apps/{0}'.format(r['name'], horsy_vars.horsypath)) + unzip('{2}apps\{0}\{1}'.format(r['name'], r['url'].split('/')[-1], horsy_vars.horsypath), + '{1}apps\{0}'.format(r['name'], horsy_vars.horsypath)) + print() - if r['download']: - if scan_to_cli('{2}apps/{0}/{1}'.format(r['name'], r['download'].split('/')[-1], - horsy_vars.horsypath))['detect']['malicious'] > 0: - print(f"[red]Dependency can be malicious. It may run now, if this added to installation " - f"config[/]") - input("Press enter if you want continue, or ctrl+c to exit") + # Scan dependencies + try: + if r['download']: + if scan_to_cli('{2}apps\{0}\{1}'.format(r['name'], r['download'].split('/')[-1], + horsy_vars.horsypath))['detect']['malicious'] > 0: + print(f"[red]Dependency can be malicious. It may run now, if this added to installation " + f"config[/]") + input("Press enter if you want continue, or ctrl+c to exit") + print() + except: + pass + # Execute install script if r['install']: print(f"Found install option") - threading.Thread(target=os.system, args=('{2}apps/{0}/{1}'.format(r['name'], r['install'], + threading.Thread(target=os.system, args=('{2}apps\{0}\{1}'.format(r['name'], r['install'], horsy_vars.horsypath),)).start() + print() - # Launch script + # Create launch script print(f"Generating launch script") - with open('{1}apps/{0}.bat'.format(r['name'], horsy_vars.horsypath), 'w') as f: + with open('{1}apps\{0}.bat'.format(r['name'], horsy_vars.horsypath), 'w') as f: f.write(f"@ECHO off\n") f.write(f"%horsypath:~0,1%:\n") f.write(f"cd %horsypath%/apps/{r['name']}\n") @@ -72,8 +89,8 @@ def install(package): except: print("[red]Unexpected error[/]") - # raise - return + raise + # return def uninstall(package): diff --git a/modules/virustotal.py b/modules/virustotal.py index 97742d5..4f4250e 100644 --- a/modules/virustotal.py +++ b/modules/virustotal.py @@ -77,5 +77,4 @@ def scan_to_cli(filename): print(f"[green]You can see report by opening: [white]{analysis['link']}[/]") print(f"{analysis['detect']['malicious']} antivirus flagged this file as malicious") - print(f"[green][OK] Done[/]") - return analysis + return analysis