def __init__(self) -> None:
self.console = Console(highlight=False)
self._crawl_progress = Progress(
TextColumn("{task.description}", table_column=Column(ratio=1)),
BarColumn(),
TimeRemainingColumn(),
expand=True,
)
self._download_progress = Progress(
TextColumn("{task.description}", table_column=Column(ratio=1)),
TransferSpeedColumn(),
DownloadColumn(),
BarColumn(),
TimeRemainingColumn(),
expand=True,
)
self._live = Live(console=self.console, transient=True)
self._update_live()
self._showing_progress = False
self._progress_suspended = False
self._lock = asyncio.Lock()
self._lines: List[str] = []
# Whether different parts of the output are enabled or disabled
self.output_explain = False
self.output_status = True
self.output_report = True
def put(
self, path: str, key: str, acl: str = "public-read", metadata: dict = {}
) -> dict:
with Progress(
SpinnerColumn(spinner_name="earth"),
TextColumn("[progress.description]{task.description}"),
BarColumn(bar_width=30),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TimeRemainingColumn(),
transient=True,
) as progress:
size = os.stat(path).st_size
task = progress.add_task(
f"[green]Uploading [bold]{os.path.basename(path)}[/bold]", total=size
)
def update_progress(complete):
progress.update(task, completed=complete)
try:
response = self.client.upload_file(
path,
self.bucket,
key,
ExtraArgs={"ACL": acl, "Metadata": metadata},
Callback=update_progress,
)
except ClientError as e:
logging.error(e)
return {}
return response
def test_progress_max_refresh() -> None:
"""Test max_refresh argment."""
time = 0.0
def get_time() -> float:
nonlocal time
try:
return time
finally:
time = time + 1.0
console = Console(color_system=None,
width=80,
legacy_windows=False,
force_terminal=True)
column = TextColumn("{task.description}")
column.max_refresh = 3
progress = Progress(
column,
get_time=get_time,
auto_refresh=False,
console=console,
)
console.begin_capture()
with progress:
task_id = progress.add_task("start")
for tick in range(6):
progress.update(task_id, description=f"tick {tick}")
progress.refresh()
result = console.end_capture()
print(repr(result))
assert (
result ==
"\x1b[?25l\r\x1b[2Kstart\r\x1b[2Kstart\r\x1b[2Ktick 1\r\x1b[2Ktick 1\r\x1b[2Ktick 3\r\x1b[2Ktick 3\r\x1b[2Ktick 5\r\x1b[2Ktick 5\n\x1b[?25h"
)
def __init__(self):
"""Init Rich Console and Table"""
super().__init__()
# init ELF table
self.table_elf = Table(title="Checksec Results: ELF", expand=True)
self.table_elf.add_column("File", justify="left", header_style="")
self.table_elf.add_column("NX", justify="center")
self.table_elf.add_column("PIE", justify="center")
self.table_elf.add_column("Canary", justify="center")
self.table_elf.add_column("Relro", justify="center")
self.table_elf.add_column("RPATH", justify="center")
self.table_elf.add_column("RUNPATH", justify="center")
self.table_elf.add_column("Symbols", justify="center")
self.table_elf.add_column("FORTIFY", justify="center")
self.table_elf.add_column("Fortified", justify="center")
self.table_elf.add_column("Fortifiable", justify="center")
self.table_elf.add_column("Fortify Score", justify="center")
# init PE table
self.table_pe = Table(title="Checksec Results: PE", expand=True)
self.table_pe.add_column("File", justify="left", header_style="")
self.table_pe.add_column("NX", justify="center")
self.table_pe.add_column("Canary", justify="center")
self.table_pe.add_column("ASLR", justify="center")
self.table_pe.add_column("Dynamic Base", justify="center")
self.table_pe.add_column("High Entropy VA", justify="center")
self.table_pe.add_column("SEH", justify="center")
self.table_pe.add_column("SafeSEH", justify="center")
self.table_pe.add_column("Force Integrity", justify="center")
self.table_pe.add_column("Control Flow Guard", justify="center")
self.table_pe.add_column("Isolation", justify="center")
# init console
self.console = Console()
# build progress bar
self.process_bar = Progress(
TextColumn("[bold blue]Processing...", justify="left"),
BarColumn(bar_width=None),
"{task.completed}/{task.total}",
"•",
"[progress.percentage]{task.percentage:>3.1f}%",
console=self.console,
)
self.display_res_bar = Progress(
BarColumn(bar_width=None),
TextColumn("[bold blue]{task.description}", justify="center"),
BarColumn(bar_width=None),
console=self.console,
transient=True,
)
self.enumerate_bar = Progress(
TextColumn("[bold blue]Enumerating...", justify="center"),
BarColumn(bar_width=None),
console=self.console,
transient=True,
)
self.process_task_id = None
def progress(self) -> Progress:
"""Creates a Progress instance.
:return: a Progress instance which can be used to display progress bars
"""
progress = Progress(TextColumn(""), BarColumn(), TextColumn("{task.percentage:0.0f}%"), console=self._console)
progress.start()
return progress
def train_one_epoch(model, optim, loss_fn, loader, epoch, steps, device,
writer, global_i, writer_interval=200, normalize=None):
model.train()
status_col = TextColumn("")
running_loss = 0
lr = optim.param_groups[0]['lr']
if normalize is not None:
assert len(normalize) == 2, "mean and std values should be provided to use data normalization"
fetcher = DataPrefetcher(loader, mean=normalize[0], std=normalize[1], device=device) # modified behavior - w/ input normalization
else:
fetcher = DataPrefetcher(loader, mean=None, std=None, device=device) # original behavior - no input normalization
samples, targets = fetcher.next()
with Progress("[progress.description]{task.description}",
"[{task.completed}/{task.total}]",
BarColumn(),
"[progress.percentage]{task.percentage:>3.0f}%",
TimeRemainingColumn(),
TextColumn("/"),
TimeElapsedColumn(),
status_col,
expand=False, console=CONSOLE, refresh_per_second=5) as progress:
task = progress.add_task(description=f'[Epoch {epoch}]', total=steps)
i = 0 # counter
t_start = time.time()
while samples is not None:
# zero the parameter gradients
optim.zero_grad()
# forward + backward + optimize
out = model(samples)
loss = loss_fn(out, targets)
loss.backward()
optim.step()
# collect running loss
running_loss += loss.item()
i += 1
global_i += 1
# update tensorboard
if i % writer_interval == 0:
writer.add_scalar('Loss/Train', running_loss/i, global_i)
# pre-fetch next samples
samples, targets = fetcher.next()
# update trackbar
if not progress.finished:
status_col.text_format = f"Loss: {running_loss/i:.06f} " \
f"speed: {(time.time() - t_start)/i:.4f}s/it " \
f"lr: {lr}"
progress.update(task, advance=1)
return running_loss / i, global_i
def __init__(self, *args, **kwargs):
columns = [
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
DownloadColumn(),
TimeRemainingColumn(),
TimeElapsedColumn()
]
kwargs['refresh_per_second'] = 1
super().__init__(*columns, *args, **kwargs)
def get_default_columns(cls) -> Tuple[ProgressColumn, ...]:
return (
TextColumn("[progress.description]{task.description}"),
MofNCompleteColumn(),
BarColumn(bar_width=30),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TimeElapsedColumn(),
TimeRemainingColumn(),
RateColumn(),
SpinnerColumn(),
)
def create_progress():
"""Returns a Rich Progress class."""
return Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}",
table_column=Column(ratio=1)),
BarColumn(bar_width=None, table_column=Column(ratio=1)),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
CustomTimeElapsedColumn(),
expand=True,
)
def track(
sequence: Union[Sequence[ProgressType], Iterable[ProgressType]],
description: str = "",
total: Optional[float] = None,
auto_refresh: bool = True,
console: Optional[Console] = None,
transient: bool = False,
get_time: Optional[Callable[[], float]] = None,
refresh_per_second: float = 10,
style: StyleType = "bar.back",
complete_style: StyleType = "bar.complete",
finished_style: StyleType = "bar.finished",
pulse_style: StyleType = "bar.pulse",
update_period: float = 0.1,
disable: bool = False,
remove: bool = False,
) -> Iterable[ProgressType]:
columns: List[ProgressColumn] = ([
TextColumn("[progress.description]{task.description}")
] if description else []) + [
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
BarColumn(
bar_width=30,
style=style,
complete_style=complete_style,
finished_style=finished_style,
pulse_style=pulse_style,
),
TimeElapsedColumn(),
TimeRemainingColumn(),
RateColumn(),
SpinnerColumn(),
]
progress = Progress(
*columns,
auto_refresh=auto_refresh,
console=console,
transient=transient,
get_time=get_time,
refresh_per_second=refresh_per_second or 10,
disable=disable,
)
with progress:
yield from progress.track(
sequence,
total=total,
description=description,
update_period=update_period,
remove=remove,
)
def __init__(self, steps: int):
super().__init__()
self.steps = steps
self.task = None
self.bar = Progress(
TextColumn(
"Iteration {task.fields[iteration]}/{task.fields[cumul]}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TextColumn("Loss = {task.fields[loss]:.5f} | ⏳ "),
TimeRemainingColumn(),
TextColumn("lr = {task.fields[lr]:.6f}"),
)
def create_bar():
return Progress(
"[progress.percentage]{task.percentage:>3.0f}%",
BarColumn(),
DownloadColumn(),
TextColumn("[red]["),
TransferSpeedColumn(),
TimeRemainingColumn(),
TextColumn("[cyan]]"),
TextColumn("[yellow]["),
TextColumn("[yellow]{task.description}", justify="center"),
TextColumn("[yellow]]"),
)
def download_all_progress_bar():
"""
Get progress bar to show progress on downloading of all source files.
"""
progress_bar = Progress(
TextColumn(
"[bold blue]Fetching files: {task.percentage:>3.0f}% ({task.completed}/{task.total})"
),
BarColumn(),
TimeElapsedColumn(),
TextColumn("({task.description})"),
)
return progress_bar
def progress(self,
prefix: str = "",
suffix: str = "{task.percentage:0.0f}%") -> Progress:
"""Creates a Progress instance.
:param prefix: the text to show before the bar (defaults to a blank string)
:param suffix: the text to show after the bar (defaults to the task's percentage)
:return: a Progress instance which can be used to display progress bars
"""
progress = Progress(TextColumn(prefix),
BarColumn(),
TextColumn(suffix),
console=self._console)
progress.start()
return progress
def parse_apk(apkfile, workers, fn_match=None, outfile=None):
console = Console()
console.log(f"Parsing {apkfile} with {workers} workers ...")
dexes = list(extract_dex_files(apkfile))
console.log(f"Found {len(dexes)} DEX file.")
total = sum(map(lambda d: len(d.data), dexes))
progress = Progress(
TextColumn("[progress.description]{task.description}"),
BarColumn(
complete_style='bar.complete',
finished_style='bar.finished',
pulse_style='bar.pulse',
),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TimeRemainingColumn(),
TimeElapsedColumn(),
console=console,
)
out = {}
out.update(JNI_COMMON)
num_classes = 0
t0 = datetime.now()
with progress:
task = progress.add_task("Analyzing...", total=total)
with multiprocessing.Pool(workers) as pool:
result = pool.imap(parse_dex_proc, dexes)
for dex, count, res in result:
if count == 0:
console.log(
"Parse {} ({} bytes) [bold red]failed: {}".format(
dex.name, len(dex.data), res))
continue
console.log("Parse {} ({} bytes), found {} classes.".format(
dex.name, len(dex.data), count))
num_classes += count
progress.update(task, advance=len(dex.data))
for cname, data in res.items():
if fn_match and not fn_match(cname):
continue
out.update(data)
console.log("Aanlyzed {} classes, cost: {}".format(num_classes,
datetime.now() - t0))
console.log("Found {} JNI methods.".format(len(out)))
if not outfile:
console.print_json(data=out)
else:
with open(outfile, 'w') as f:
json.dump(out, f, indent=2, ensure_ascii=False)
def check_remote_servers(logger, servers):
check_result = dict()
progress = Progress(
"{task.description}",
SpinnerColumn(),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
)
with progress:
with ThreadPoolExecutor() as pool:
for server in servers:
if server["check"] == "False":
continue
reader = DriverReader(progress)
tinfo = pool.submit(
reader.get_remote_drivers,
server["ip"],
server["user"],
server["password"],
server["ssh_port"],
server["query"],
)
check_result[server["ip"]] = tinfo
for ip in check_result:
drivers, wu_drivers, noinfo_drivers = check_result[ip].result()
check_result[ip] = {
"drivers": drivers,
"weak-update-drivers": wu_drivers,
"noinfo-drivers": noinfo_drivers,
}
progress.console.print("[green]Check completed![/]")
return check_result
def download(self, batch, articles, folder):
"""
Download a pdf batch
"""
log_file = os.path.join(folder, 'missing.log')
logger.remove()
logger.add(log_file,
format="{time} {level} {message}",
mode='w',
level="INFO")
assert len(articles) > 0, 'no article.'
progress = Progress(TextColumn(
"[progress.description]{task.description}",
table_column=Column(ratio=1)),
TimeElapsedColumn(table_column=Column(ratio=1)),
BarColumn(table_column=Column(ratio=2)),
"| {task.completed} of {task.total:>2.0f}",
expand=False)
missing = 0
with progress:
task = progress.add_task(f" {batch} |", total=len(articles))
for article in articles:
done = self.get_article(
article['article_url'],
os.path.join(folder, article['file_name']))
if done:
progress.update(task, advance=1)
else:
missing += 1
logger.info("NOT_FOUND_IN_SCI-HUB | " +
article['warning_str'])
return missing, log_file
def _create_rich_progress(self):
desc_column = TextColumn("{task.description}", table_column=Column(ratio=1))
bar_column = BarColumn(bar_width=None, table_column=Column(ratio=2))
progress = Progress(desc_column, bar_column, "[progress.percentage]{task.percentage}%", expand=True)
return progress
def dump_mem(vm_name: Optional[str], kvm_unix_socket: Optional[str], output_file: Optional[str]):
# prepare drivers init params
init_params = DriverInitParamsPy()
common = CommonInitParamsPy()
common.vm_name = vm_name
init_params.common = common
if kvm_unix_socket:
kvm = KVMInitParamsPy()
kvm.unix_socket = kvm_unix_socket
init_params.kvm = kvm
# init libmicrovmi
micro = Microvmi(None, init_params)
destination = output_file.format(vm_name=vm_name)
with pause_ctxt(micro):
max_addr = micro.max_addr
print(f"Dumping physical memory on {vm_name} until 0x{max_addr:X} to {destination}")
with open(destination, "wb") as f:
with Progress(
"[progress.description]{task.description}",
TextColumn("[bold yellow]0x{task.completed:X}"),
BarColumn(bar_width=None),
"[progress.percentage]{task.percentage:>3.0f}%",
transient=True,
) as progress:
dump_task = progress.add_task("Dumping ", total=max_addr)
mem = micro.padded_memory
for addr in range(0, max_addr, READ_SIZE):
logging.debug("dumping at 0x%x", addr)
current_chunk_size = min(READ_SIZE, max_addr - addr)
buffer = mem.read(current_chunk_size)
f.write(buffer)
progress.update(dump_task, advance=READ_SIZE)
def create_progress_bar():
"""
Create progress bar to display overall progress.
Returns rich.progress.Progress instance if the Rich Python package is available,
or a shim DummyProgress instance otherwise.
"""
if use_rich() and build_option('show_progress_bar'):
# pick random spinner, from a selected subset of available spinner (see 'python3 -m rich.spinner')
spinner = random.choice(('aesthetic', 'arc', 'bounce', 'dots', 'line',
'monkey', 'point', 'simpleDots'))
progress_bar = Progress(
SpinnerColumn(spinner),
"[progress.percentage]{task.percentage:>3.1f}%",
TextColumn(
"[bold blue]Installing {task.description} ({task.completed:.0f}/{task.total} done)"
),
BarColumn(bar_width=None),
TimeElapsedColumn(),
transient=True,
expand=True,
)
else:
progress_bar = DummyProgress()
return progress_bar
def UploadFile(socket, address, key, size_uncompressed, size_compressed, buffer=2048):
with open("temp.tar.gz", "rb") as f:
file_hash_uc = SHA512.new();
file_hash_c = SHA512.new();
for address_singular in address:
with open(address_singular, "rb") as filehandle:
while True:
block = filehandle.read(buffer);
if not block:
break;
file_hash_uc.update(block);
with Progress(TextColumn("[bold blue]{task.description}", justify="right"),
BarColumn(bar_width=None),
"[progress.percentage]{task.percentage:>3.1f}%",
"•",
DownloadColumn(),
"•",
TransferSpeedColumn(),
"•",
TimeRemainingColumn(),) as progress:
task = progress.add_task("Uploading file(s)", total=size_compressed);
while not progress.finished:
l = f.read(buffer);
if not l:
break;
select.select([], [socket], []);
sendEncryptedMessage(socket, l, key);
progress.update(task, advance=len(l));
file_hash_c.update(l);
return (file_hash_uc, file_hash_c);
def __init__(
self,
console: Console,
flow: RunningBatchFlow,
bake: Bake,
attempt: Attempt,
client: RetryReadNeuroClient,
storage: AttemptStorage,
bake_storage: BakeStorage,
project_storage: ProjectStorage,
*,
polling_timeout: Optional[float] = 1,
project_role: Optional[str] = None,
) -> None:
self._progress = Progress(
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TextColumn("[progress.remaining]{task.elapsed:.0f} sec"),
console=console,
auto_refresh=False,
transient=self.transient_progress,
redirect_stderr=True,
)
self._top_flow = flow
self._bake = bake
self._attempt = attempt
self._client = client
self._storage = storage
self._project_storage = project_storage
self._bake_storage = bake_storage
self._graphs: Graph[RunningBatchBase[BaseBatchContext]] = Graph(
flow.graph, flow)
self._tasks_mgr = BakeTasksManager()
self._is_cancelling = False
self._project_role = project_role
self._is_projet_role_created = False
self._run_builder_job = self.run_builder_job
# A note about default value:
# AS: I have no idea what timeout is better;
# too short value bombards servers,
# too long timeout makes the waiting longer than expected
# The missing events subsystem would be great for this task :)
self._polling_timeout = polling_timeout
self._bars: Dict[FullID, TaskID] = {}
def configure_columns(self, trainer) -> list:
return [
TextColumn("[progress.description]{task.description}"),
BarColumn(),
BatchesProcessedColumn(),
TimeRemainingColumn(),
TimeElapsedColumn()
]
def progress(self) -> Progress:
"""
Return a pre-configured rich spinner.
"""
text_column = TextColumn("{task.description}")
spinner_column = SpinnerColumn("simpleDotsScrolling",
style="bold white")
return Progress(text_column, spinner_column, transient=True)
def get_metric_text_column(self):
return TextColumn(
"\t".join(
f"{key}: {{task.fields[{key}]}}"
for key, value in self.metrics.items()
if (key in self.metrics_to_receive) or (key in self.metrics_to_track)
).expandtabs(2)
)
def evaluate(model, loss_fn, loader, epoch, steps, device, normalize=None):
model.eval()
status_col = TextColumn("")
running_loss = 0
if normalize is not None:
assert len(normalize) == 2, "mean and std values should be provided to use data normalization"
fetcher = DataPrefetcher(loader, mean=normalize[0], std=normalize[1], device=device) # modified behavior - w/ input normalization
else:
fetcher = DataPrefetcher(loader, mean=None, std=None, device=device) # original behavior - no input normalization
samples, targets = fetcher.next()
with Progress("[progress.description]{task.description}",
"[{task.completed}/{task.total}]",
BarColumn(),
"[progress.percentage]{task.percentage:>3.0f}%",
TimeRemainingColumn(),
TextColumn("/"),
TimeElapsedColumn(),
status_col,
expand=False, console=CONSOLE, refresh_per_second=5) as progress:
task = progress.add_task(description=f'[Eval {epoch}]', total=steps)
i = 0 # counter
t_start = time.time()
with torch.no_grad():
while samples is not None:
# forward only
out = model(samples)
val_loss = loss_fn(out, targets)
# collect running loss
running_loss += val_loss.item()
i += 1
# pre-fetch next samples
samples, targets = fetcher.next()
if not progress.finished:
status_col.text_format = f"Val loss: {running_loss/i:.06f} " \
f"speed: {(time.time() - t_start)/i:.4f}s/it"
progress.update(task, advance=1)
return running_loss / i
def main():
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
transient=True,
) as progress:
progress.add_task(description="Processing...", total=None)
progress.add_task(description="Preparing...", total=None)
time.sleep(5)
print("Done!")
def rich_bar():
return Progress(
TextColumn("[bold blue]{task.description}", justify="right"),
"[progress.percentage]{task.percentage:>3.0f}%",
BarColumn(bar_width=None),
TextColumn(
"[bold blue][progress.completed]{task.completed}/{task.total}",
justify="right",
),
"[",
TimeElapsedColumn(),
"<",
TimeRemainingColumn(),
"|",
SpeedColumn(".1f"),
TextColumn("[progress.data.speed]{task.fields[unit]}/s",
justify="right"),
"]",
auto_refresh=False,
)
def extensions_progress_bar():
"""
Get progress bar to show progress for installing extensions.
"""
progress_bar = Progress(
TextColumn("[bold blue]{task.description}"),
BarColumn(),
TimeElapsedColumn(),
)
return progress_bar
def download_one_progress_bar_unknown_size():
"""
Get progress bar to show progress for downloading a file of unknown size.
"""
progress_bar = Progress(
TextColumn('[bold yellow]Downloading {task.description}'),
FileSizeColumn(),
TransferSpeedColumn(),
)
return progress_bar
