def make_progress() -> Progress:
_time = 0.0
def fake_time():
nonlocal _time
try:
return _time
finally:
_time += 1
console = Console(
file=io.StringIO(),
force_terminal=True,
color_system="truecolor",
width=80,
legacy_windows=False,
)
progress = Progress(console=console, get_time=fake_time, auto_refresh=False)
task1 = progress.add_task("foo")
task2 = progress.add_task("bar", total=30)
progress.advance(task2, 16)
task3 = progress.add_task("baz", visible=False)
task4 = progress.add_task("egg")
progress.remove_task(task4)
task4 = progress.add_task("foo2", completed=50, start=False)
progress.stop_task(task4)
progress.start_task(task4)
progress.update(
task4, total=200, advance=50, completed=200, visible=True, refresh=True
)
progress.stop_task(task4)
return progress
def run(self, manager: "pwncat.manager.Manager", args):
# Create a progress bar for the download
progress = Progress(
TextColumn("[bold cyan]{task.fields[filename]}", justify="right"),
BarColumn(bar_width=None),
"[progress.percentage]{task.percentage:>3.1f}%",
"•",
DownloadColumn(),
"•",
TransferSpeedColumn(),
"•",
TimeRemainingColumn(),
)
if not args.destination:
args.destination = f"./{os.path.basename(args.source)}"
try:
length = os.path.getsize(args.source)
started = time.time()
with progress:
task_id = progress.add_task("upload",
filename=args.destination,
total=length,
start=False)
with open(args.source, "rb") as source:
with manager.target.platform.open(args.destination,
"wb") as destination:
progress.start_task(task_id)
copyfileobj(
source,
destination,
lambda count: progress.update(task_id,
advance=count),
)
progress.update(task_id,
filename="draining buffers...")
progress.stop_task(task_id)
progress.start_task(task_id)
progress.update(task_id, filename=args.destination)
elapsed = time.time() - started
console.log(f"uploaded [cyan]{human_readable_size(length)}[/cyan] "
f"in [green]{human_readable_delta(elapsed)}[/green]")
except (FileNotFoundError, PermissionError, IsADirectoryError) as exc:
self.parser.error(str(exc))
def event(jid="all", tag=None, progress="log", stop_signal=None):
"""
Function to listen to events emitted by Nornir Proxy Minions. Matched
event printed to terminal.
:param tag: (str) tag regex string, default is ``nornir\-proxy/.*``
:param jid: (int, str) Job ID to listen events for, default is ``all``
:param progress: (str) progress display mode - ``log``, ``raw``, ``bars``, ``tree``
:param stop_signal: (obj) thread Event object, stops listening to events if ``stop_signal.is_set()``,
if ``stop_signal is None``, listens and print events until keyboard interrupt hit - ``ctrl+c``
``bars`` and ``tree`` progress display modes use Rich library, to properly display various
symbols and characters need to make sure to use utf-8 encoding for your environment for example
by running these commands::
[root@salt-master ~]# PYTHONIOENCODING=utf-8
[root@salt-master ~]# export PYTHONIOENCODING
"""
# start rich console
globals().setdefault("console", Console())
stop_signal = stop_signal or Event()
events_queue = queue.Queue()
tag = (
tag
if tag
else (
r"nornir\-proxy\/.*"
if jid == "all"
else r"nornir\-proxy\/{jid}\/.*".format(jid=jid)
)
)
# start events thread
listen_events_thread = Thread(
target=_get_salt_nornir_event,
kwargs={"stop_signal": stop_signal, "tag": tag, "events_queue": events_queue},
daemon=True, # to not block once main process finishes
)
listen_events_thread.start()
# display events
if HAS_RICH and progress == "bars":
tasks = {}
stop_events_loop = Event()
rich_progress = Progress(
TextColumn("{task.description}"),
BarColumn(),
TextColumn("{task.percentage:>3.0f}%"),
TextColumn("{task.completed}/{task.total}"),
TimeElapsedColumn(),
TextColumn("{task.fields[info]}"),
TextColumn("{task.fields[status]}"),
refresh_per_second=5,
)
# listen to events indefinitely if stop_signal is None
with rich_progress:
while True:
try:
e = events_queue.get(timeout=1)
except queue.Empty:
continue
finally:
# check if need to finish
if stop_signal.is_set() and all(
[t.finished for t in rich_progress.tasks]
):
stop_events_loop.set()
break
edata = e["data"]
jid = edata["jid"]
task_type = edata["task_type"]
task_event = edata["task_event"]
status = edata["status"]
task_name = edata["task_name"]
status_str = (
f"[red]{status}" if status == "FAILED" else f"[green]{status}"
)
description_str = (
"[blue]{timestamp}[/]:{user}:{jid}:[bright_magenta]{function}[/]"
)
# catch task started events and add new progress bar
if task_type == "task" and task_event == "started":
# if job runs on multiple proxy minions
if jid in tasks:
task = rich_progress.tasks[tasks[jid]]
task.total += len(edata["hosts"])
rich_progress.update(tasks[jid])
else:
total = len(edata["hosts"])
description = description_str.format(**edata)
info = "[cyan]{}[/]:{}".format(
task_type, task_name.split(".")[-1]
)
tasks[jid] = rich_progress.add_task(
description, total=total, info=info, status=status
)
# catch task instance end events to advance progress bar
elif task_type == "task_instance" and task_event == "completed":
if jid in tasks:
task = rich_progress.tasks[tasks[jid]]
if status == "PASSED":
rich_progress.update(tasks[jid], advance=1)
if task.completed >= task.total or status == "FAILED":
rich_progress.tasks[tasks[jid]].fields[
"status"
] = status_str
elif task_type == "task" and task_event == "completed":
if jid in tasks:
task = rich_progress.tasks[tasks[jid]]
if task.completed >= task.total or status == "FAILED":
rich_progress.tasks[tasks[jid]].fields[
"status"
] = status_str
rich_progress.stop_task(tasks[jid])
# stop all subtasks for this jid and update their status
for task in tasks:
if task.startswith(jid):
rich_progress.tasks[tasks[task]].fields[
"status"
] = status_str
rich_progress.stop_task(tasks[task])
# handle subtask progress
elif task_type == "subtask" and task_event == "started":
tid = "{jid}:{task_name}".format(**edata)
if tid not in tasks and jid in tasks:
total = rich_progress.tasks[tasks[jid]].total
description = description_str.format(**edata)
info = "[cyan]{task_type}[/]:{task_name}".format(**edata)
tasks[tid] = rich_progress.add_task(
description, total=total, info=info, status=status
)
# update task total if got additional start events
elif tid in tasks and jid in tasks:
task = rich_progress.tasks[tasks[tid]]
task.total = rich_progress.tasks[tasks[jid]].total
rich_progress.update(tasks[tid])
elif task_type == "subtask" and task_event == "completed":
tid = "{jid}:{task_name}".format(**edata)
if tid in tasks:
task = rich_progress.tasks[tasks[tid]]
if status == "PASSED":
rich_progress.update(tasks[tid], advance=1)
if task.completed >= task.total or status == "FAILED":
rich_progress.tasks[tasks[tid]].fields[
"status"
] = status_str
# hide tasks beyond vertical overflow point
hght = console.size.height - 2
for tindx, _ in enumerate(rich_progress.tasks[:-hght]):
rich_progress.update(tindx, visible=False)
elif progress == "tree":
pass
elif progress == "raw":
while not stop_signal.is_set():
try:
e = events_queue.get(timeout=1)
print(e)
except queue.Empty:
continue
# handle 'log' progress mode
else:
while not stop_signal.is_set():
try:
e = events_queue.get(timeout=1)
except queue.Empty:
continue
edata = e["data"]
# form message string and print it
if edata["task_type"] == "task":
edata["hs"] = ", ".join(edata["hosts"])
elif edata["task_type"] in ["task_instance", "subtask"]:
edata["hs"] = edata["host"]
msg = "{timestamp}:{user}:{jid}:{proxy_id}:w{worker_id} {function} {hs} {task_type} {task_event} {task_name}; {status}"
print(msg.format(**edata))
class RichProgressBar(ProgressBarBase):
def __init__(self):
super().__init__()
self.pg = Progress(
"[progress.description]{task.description}",
BarColumn(),
"[progress.percentage]{task.completed}/{task.total}",
BarDictColumn(),
transient=True,
expand=True,
refresh_per_second=2,
)
self._train_id = self._val_id = self._test_id = None
def disable(self):
self.pg.disable = True
def enable(self):
self.pg.disable = False
def on_train_start(self, trainer, pl_module):
super().on_train_start(trainer, pl_module)
self.pg.start()
def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx,
dataloader_idx):
super().on_test_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
bardic = pl_module.get_progress_bar_dict()
self.pg.update(self._test_id, completed=self.test_batch_idx, **bardic)
if self.test_batch_idx >= self.total_test_batches:
self.pg.stop_task(self._test_id)
def on_test_epoch_start(self, trainer, pl_module) -> None:
super().on_test_epoch_start(trainer, pl_module)
if self._test_id is not None:
self.pg.remove_task(self._test_id)
self._test_id = self.pg.add_task("epoch T%03d" % trainer.current_epoch,
total=self.total_test_batches)
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx,
dataloader_idx):
super().on_train_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
bardic = pl_module.get_progress_bar_dict()
self.pg.update(self._train_id,
completed=self.train_batch_idx,
**bardic)
if self.train_batch_idx >= self.total_train_batches:
self.pg.stop_task(self._train_id)
def on_train_epoch_start(self, trainer, pl_module):
super().on_train_epoch_start(trainer, pl_module)
if self._train_id is not None:
self.pg.remove_task(self._train_id)
if self._val_id is not None:
self.pg.remove_task(self._val_id)
self._train_id = self.pg.add_task("epoch T%03d" %
trainer.current_epoch,
total=self.total_train_batches)
def on_validation_batch_end(self, trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx):
super().on_validation_batch_end(trainer, pl_module, outputs, batch,
batch_idx, dataloader_idx)
bardic = pl_module.get_progress_bar_dict()
self.pg.update(self._val_id, completed=self.val_batch_idx, **bardic)
if self.val_batch_idx >= self.total_val_batches:
self.pg.stop_task(self._val_id)
def on_validation_epoch_start(self, trainer, pl_module) -> None:
super().on_validation_epoch_start(trainer, pl_module)
self._val_id = self.pg.add_task("epoch V%03d" % trainer.current_epoch,
total=self.total_val_batches)
def on_train_end(self, trainer, pl_module):
super().on_train_end(trainer, pl_module)
self.pg.stop()
def print(self, *args, **kwargs):
self.pg.console.print(*args, **kwargs)
class RichLogger(object):
"""Defines a logger based on `rich.RichHandler`.
Compared to the basic Logger, this logger will decorate the log message in
a pretty format automatically.
"""
def __init__(self,
work_dir=DEFAULT_WORK_DIR,
logfile_name='log.txt',
logger_name='logger'):
"""Initializes the logger.
Args:
work_dir: The work directory. (default: DEFAULT_WORK_DIR)
logfile_name: Name of the log file. (default: `log.txt`)
logger_name: Unique name for the logger. (default: `logger`)
"""
self.logger = logging.getLogger(logger_name)
if self.logger.hasHandlers(): # Already existed
raise SystemExit(
f'Logger `{logger_name}` has already existed!\n'
f'Please use another name, or otherwise the '
f'messages from these two logger may be mixed up.')
self.logger.setLevel(logging.DEBUG)
# Print log message with `INFO` level or above onto the screen.
terminal_console = Console(file=sys.stderr,
log_time=False,
log_path=False)
terminal_handler = RichHandler(level=logging.INFO,
console=terminal_console,
show_time=True,
show_level=True,
show_path=False)
terminal_handler.setFormatter(logging.Formatter('%(message)s'))
self.logger.addHandler(terminal_handler)
# Save log message with all levels into log file if needed.
if logfile_name:
os.makedirs(work_dir, exist_ok=True)
file_stream = open(os.path.join(work_dir, logfile_name), 'a')
file_console = Console(file=file_stream,
log_time=False,
log_path=False)
file_handler = RichHandler(level=logging.DEBUG,
console=file_console,
show_time=True,
show_level=True,
show_path=False)
file_handler.setFormatter(logging.Formatter('%(message)s'))
self.logger.addHandler(file_handler)
self.log = self.logger.log
self.debug = self.logger.debug
self.info = self.logger.info
self.warning = self.logger.warning
self.error = self.logger.error
self.exception = self.logger.exception
self.critical = self.logger.critical
self.pbar = None
def print(self, *messages, **kwargs):
"""Prints messages without time stamp or log level."""
for handler in self.logger.handlers:
handler.console.print(*messages, **kwargs)
def init_pbar(self, leave=False):
"""Initializes a progress bar which will display on the screen only.
Args:
leave: Whether to leave the trace. (default: False)
"""
assert self.pbar is None
# Columns shown in the progress bar.
columns = (
TextColumn("[progress.description]{task.description}"),
BarColumn(bar_width=None),
TextColumn("[progress.percentage]{task.percentage:>5.1f}%"),
TimeColumn(),
)
self.pbar = Progress(*columns,
console=self.logger.handlers[0].console,
transient=not leave,
auto_refresh=True,
refresh_per_second=10)
self.pbar.start()
def add_pbar_task(self, name, total):
"""Adds a task to the progress bar.
Args:
name: Name of the new task.
total: Total number of steps (samples) contained in the task.
Returns:
The task ID.
"""
assert isinstance(self.pbar, Progress)
task_id = self.pbar.add_task(name, total=total)
return task_id
def update_pbar(self, task_id, advance=1):
"""Updates a certain task in the progress bar.
Args:
task_id: ID of the task to update.
advance: Number of steps advanced onto the target task. (default: 1)
"""
assert isinstance(self.pbar, Progress)
if self.pbar.tasks[int(task_id)].finished:
if self.pbar.tasks[int(task_id)].stop_time is None:
self.pbar.stop_task(task_id)
else:
self.pbar.update(task_id, advance=advance)
def close_pbar(self):
"""Closes the progress bar"""
assert isinstance(self.pbar, Progress)
self.pbar.stop()
self.pbar = None
def train(self, trainData, trainLabel):
'''
训练
===
Arguments
---------
- `trainData` 训练集数据
- `trainLabel` 训练集标签
Algorithm
---------
- Sequential minimal optimization, SMO
Returns
-------
'''
self.__x, self.__y = np.array(trainData), np.array(trainLabel)
self.__alpha = np.zeros(self.__x.shape[0])
self.__b = 0
self.__K = np.zeros([self.__x.shape[0], self.__x.shape[0]],
dtype=float) # 训练数据核函数表
for i in range(self.__x.shape[0]):
for j in range(i, self.__x.shape[0]):
self.__K[i, j] = self.__K[j, i] = self.Kernel(
self.__x[i], self.__x[j]) # 计算核函数表
progress = Progress(
"[progress.description]{task.description}",
BarColumn(bar_width=None),
"[progress.percentage]{task.completed}/{task.total}",
"•",
"[time]{task.elapsed:.2f}s",
) # rich 进度条
progress.start()
allSatisfied = False # 全部满足 KKT 条件
iteration = 1 # 迭代次数
while not allSatisfied:
allSatisfied = True
iterateTask = progress.add_task(
"[yellow]{} iterating...".format(iteration),
total=self.__x.shape[0])
iteration += 1
for i in range(self.__x.shape[0]): # 外层循环
progress.update(iterateTask, advance=1)
if not (self.__ifSatisfyKKT(i)): # 选择第一个变量
E1 = self.__Error(i)
maximum = -1
for k in range(self.__x.shape[0]): # 内层循环
tempE = self.__Error(k)
tempE_difference = np.fabs(E1 - self.__Error(k))
if tempE_difference > maximum: # 选择第二个变量
maximum = tempE_difference
E2 = tempE
j = k
if maximum == -1:
continue
U = max(0, (self.__alpha[i] + self.__alpha[j] -
self.__C) if self.__y[i] == self.__y[j] else
(self.__alpha[j] -
self.__alpha[i])) # alpha^2_new 的下界
V = min(
self.__C,
(self.__alpha[i] +
self.__alpha[j]) if self.__y[i] == self.__y[j] else
(self.__alpha[j] - self.__alpha[i] +
self.__C)) # alpha^new_2 的上界
alpha_2_new = self.__alpha[j] + self.__y[j] * (E1 - E2) / (
self.__K[i, i] + self.__K[j, j] - 2 * self.__K[i, j])
# alpha^2_new 越界
if alpha_2_new > V:
alpha_2_new = V
elif alpha_2_new < U:
alpha_2_new = U
alpha_1_new = self.__alpha[i] + self.__y[i] * \
self.__y[j] * (self.__alpha[j] - alpha_2_new)
# 更新偏置
b_1_new = -E1 - self.__y[i] * self.__K[i, i] * (
alpha_1_new -
self.__alpha[i]) - self.__y[j] * self.__K[j, i] * (
alpha_2_new - self.__alpha[j]) + self.__b
b_2_new = -E2 - self.__y[i] * self.__K[i, j] * (
alpha_1_new -
self.__alpha[i]) - self.__y[j] * self.__K[j, j] * (
alpha_2_new - self.__alpha[j]) + self.__b
# 实装更新
if (np.fabs(self.__alpha[i] - alpha_1_new) <
0.0000001) and (np.fabs(self.__alpha[j] -
alpha_2_new) < 0.0000001):
continue
else:
allSatisfied = False
self.__alpha[i] = alpha_1_new
self.__alpha[j] = alpha_2_new
if 0 < alpha_1_new < self.__C:
self.__b = b_1_new
elif 0 < alpha_2_new < self.__C:
self.__b = b_2_new
else:
self.__b = (b_1_new + b_2_new) / 2
progress.stop_task(iterateTask)
progress.stop()
return
with Live(progress_group):
for idx, (name, step_times) in enumerate(apps):
# update message on overall progress bar
top_descr = "[bold #AAAAAA](%d out of %d apps installed)" % (idx,
len(apps))
overall_progress.update(overall_task_id, description=top_descr)
# add progress bar for steps of this app, and run the steps
current_task_id = current_app_progress.add_task("Installing app %s" %
name)
app_steps_task_id = app_steps_progress.add_task("",
total=len(step_times),
name=name)
run_steps(name, step_times, app_steps_task_id)
# stop and hide steps progress bar for this specific app
app_steps_progress.update(app_steps_task_id, visible=False)
current_app_progress.stop_task(current_task_id)
current_app_progress.update(
current_task_id,
description="[bold green]App %s installed!" % name)
# increase overall progress now this task is done
overall_progress.update(overall_task_id, advance=1)
# final update for message on overall progress bar
overall_progress.update(
overall_task_id,
description="[bold green]%s apps installed, done!" % len(apps))
