def handle(self, *args, **options):
name = options["metric"]
try:
metric = metrics_registry[name]()
except KeyError:
self.console.print(Text(f"Metric '{name}' not found!", style="bold red"))
raise CommandError(1)
results = metric.get_data()
for result in results:
if result.holds_data:
table = Table(
title=f"{escape(result.alias)} ({escape(result.dsn)})",
title_style="bold green",
)
for header in metric.headers:
table.add_column(escape(header.name), no_wrap=True)
for record in result.records:
table.add_row(
*[
Text(
str(item),
style=RICH_STYLE_MAPPING.get(
metric.get_record_item_style(record, item, idx)
),
)
for idx, item in enumerate(record)
],
style=RICH_STYLE_MAPPING.get(metric.get_record_style(record)),
)
self.console.print(table)
else:
self.console.print(escape(result.reason), style="bold red")
def error_exit(reason, exit_status=1):
errcon.print(escape(program_name),
error_marker,
" ",
escape(reason),
sep='')
sys.exit(exit_status)
def test_escape():
# Potential tags
assert escape("foo[bar]") == r"foo\[bar]"
assert escape(r"foo\[bar]") == r"foo\\\[bar]"
# Not tags (escape not required)
assert escape("[5]") == "[5]"
assert escape("\\[5]") == "\\[5]"
def test_render_escape():
console = Console(width=80, color_system=None)
console.begin_capture()
console.print(escape(r"[red]"), escape(r"\[red]"), escape(r"\\[red]"),
escape(r"\\\[red]"))
result = console.end_capture()
expected = r"[red] \[red] \\[red] \\\[red]" + "\n"
assert result == expected
def remove_channel(self, name):
'''
Attempt to remove a channel, printing progress output as we go and
returning a boolean that indicates whether we think we removed the
channel.
'''
if name.startswith('release'):
print('[red]This tool does not mess with release channels![/red]')
return HelperResult(HANDS_OFF_RELEASE_CHANNELS)
all_channels = self.determine_channels()
channel = all_channels.get(name, None)
if channel is None:
print('[red]No such channel[/red]:', name)
return HelperResult(NO_SUCH_CHANNEL)
lbArn = channel.get('loadBalancerArn', None)
if lbArn:
print('Located channel load balancer ARN:', escape(lbArn))
# Per docs this also deletes the associated listeners (and their
# rules).
r = self.elb.delete_load_balancer(LoadBalancerArn=lbArn)
print('Deleted load balancer (and listeners)!', r)
else:
print('[yellow]There was no load balancer to delete[/yellow]')
tgArn = channel.get('targetGroupArn', None)
if tgArn:
print('Located channel target group ARN:', escape(tgArn))
r = self.elb.delete_target_group(TargetGroupArn=tgArn)
print('Deleted target group!', r)
else:
print('[yellow]There was no target group to delete[/yellow]')
rrset = channel.get('rrset')
if rrset:
print('[green]Located channel subdomain:[/green]', rrset)
r = self.dns.change_resource_record_sets(
HostedZoneId=channel['zoneId'],
ChangeBatch={
'Comment':
f"remove channel {name}",
'Changes': [{
'Action': 'DELETE',
'ResourceRecordSet': rrset,
}],
})
print('[green]Deleted subdomain:[/green]', r)
else:
print('[yellow]There was no DNS sub-domain to delete[/yellow]')
return HelperResult()
def print_location_error(error, line=None, file=sys.stderr):
errcon.print(escape(str(error.location)), style='bold', end='')
errcon.print(error_marker, escape(error.message_without_location))
if line is not None:
errcon.print(
rich.syntax.Syntax(line, "omg-idl",
start_line=error.location.line))
errcon.print(' ' * (error.location.col - 1),
'^',
'~' * (error.location.length - 1),
style='bold red',
sep='')
def handle(self, *args, **options):
table = Table()
table.add_column("Slug", no_wrap=True)
table.add_column("Label", no_wrap=True)
table.add_column("Description")
for metric in metrics_registry.sorted:
description = metric.__doc__ or ""
table.add_row(
escape(str(metric.slug)),
escape(str(metric.label)),
escape(textwrap.dedent(description).strip()),
)
self.console.print(table)
def run(s: Optional[str] = None, race: bool = True) -> bool:
success = True
env = os.environ.copy()
env["DEBUG"] = "true"
if s is None:
p = subprocess.Popen(
['go', 'test', '-v', '-count=1', '-race' if race else ''],
stdout=subprocess.PIPE,
env=env)
else:
p = subprocess.Popen([
'go', 'test', '-v', '-run', s, '-count=1', '-race' if race else ''
],
stdout=subprocess.PIPE,
env=env)
output = ""
for line in iter(p.stdout.readline, b''):
out = line.decode('utf-8')
output += out
out = out.strip("\n")
if "INFO" in out:
continue
if "PASS" in out:
print(f"[green]{escape(out)}[/green]")
elif "FAIL" in out:
print(f"[red]{escape(out)}[/red]")
success = False
else:
print(escape(out))
if not success:
fn = f"{s + '-' if s is not None else ''}fail-{random.randint(1,10000)}"
print(f"[magenta]saving failed log file to {fn}")
with open(fn, "w") as f:
f.write(output)
return success
def print(self):
if not self:
print()
print("[red][No Rename][/red]")
print()
return
print()
table = Table(box=box.ROUNDED)
table.add_column("old_name", justify="right")
table.add_column("new_name", justify="left")
for file in self.files:
_old_ = escape(file.old_name)
_new_ = escape(file.new_name)
table.add_row(_old_, _new_)
print(table)
print()
def status(self, style: str, action: str, text: str) -> None:
"""
Print a status update while crawling. Allows markup in the "style"
argument which will be applied to the "action" string.
"""
if self.output_status:
action = escape(f"{action:<{self.STATUS_WIDTH}}")
self.print(f"{style}{action}[/] {escape(text)}")
def prettify(gadget):
# things like dword ptr [eax] need to be escaped for rich's markup
escaped = escape(str(gadget))
escaped = re.sub(r"(0x[0-9a-fA-F]+(?!\]))", r"[blue]\1[/]", escaped)
escaped = escaped.replace(f"ret", f"[red]ret[/]")
for instr in INSTRUCTIONS:
escaped = escaped.replace(f" {instr}", f" [steel_blue1]{instr}[/]")
return escaped
def _lines() -> Iterator[str]:
yield (f"The command [command]git {event.command}[/] failed"
f" with status [status]{event.status}[/].")
command = escape(shlex.join(["git", event.command,
*event.options]))
yield ""
yield f"❯ {command}"
for stream in ["stdout", "stderr"]:
output = getattr(event, stream)
if output:
yield ""
for line in output.splitlines():
line = escape(line)
yield f"[{stream}]{line}[/]"
yield ""
def describe(self, t=None):
if not t:
t = Tree("")
if self.configuration_key:
sub_t = t.add(
f"[deep_sky_blue1]configuration[/deep_sky_blue1]: "
f"[orange3]{self.configuration_key}"
)
if self.__doc__:
t.add(f"[deep_sky_blue1]doc[/deep_sky_blue1]: {self.__doc__.strip()}")
if self._item_type and self._return_type:
sub_t = t.add(
f"[deep_sky_blue1]signature[/deep_sky_blue1]: "
f"{pretty_print_type(self._item_type)} ->"
f" {pretty_print_type(self._return_type)}"
)
if self._load_time:
sub_t = t.add(
"[deep_sky_blue1]load time[/deep_sky_blue1]: [orange3]"
+ humanize.naturaldelta(self._load_time, minimum_unit="seconds")
)
if self._load_memory_increment is not None:
sub_t = t.add(
f"[deep_sky_blue1]load memory[/deep_sky_blue1]: "
f"[orange3]{humanize.naturalsize(self._load_memory_increment)}"
)
if self.model_dependencies.models:
dep_t = t.add("[deep_sky_blue1]dependencies")
for m in self.model_dependencies.models:
dep_t.add("[orange3]" + escape(m))
if self.asset_path:
sub_t = t.add(
f"[deep_sky_blue1]asset path[/deep_sky_blue1]: "
f"[orange3]{self.asset_path}"
)
if self.batch_size:
sub_t = t.add(
f"[deep_sky_blue1]batch size[/deep_sky_blue1]: "
f"[orange3]{self.batch_size}"
)
if self.model_settings:
sub_t = t.add("[deep_sky_blue1]model settings[/deep_sky_blue1]")
describe(self.model_settings, t=sub_t)
return t
def download_bar(
self,
style: str,
action: str,
text: str,
total: Optional[float] = None,
) -> ContextManager[ProgressBar]:
"""
Allows markup in the "style" argument which will be applied to the
"action" string.
"""
action = escape(f"{action:<{self.STATUS_WIDTH}}")
description = f"{style}{action}[/] {text}"
return self._bar(self._download_progress, description, total)
def __construct_file_tree(folder: str, basename: str) -> Tuple[FileTree, int, int]:
"""
Recurses through the project directories.
Constructs a file tree by subsequent calls to the API
"""
tree = FileTree([], f"{basename}/")
try:
sorted_files_folders = __api_call_list_files(folder)
except exceptions.NoDataError as e:
if folder is None:
raise exceptions.NoDataError(
"No files or folders found for the specified project"
)
else:
raise exceptions.NoDataError(f"Could not find folder: '{escape(folder)}'")
# Get max length of file name
max_string = max([len(x["name"]) for x in sorted_files_folders])
# Get max length of size string
max_size = max(
[
len(x["size"].split(" ")[0])
for x in sorted_files_folders
if show_size and "size" in x
],
default=0,
)
# Rich outputs precisely one line per file/folder
for f in sorted_files_folders:
is_folder = f.pop("folder")
if not is_folder:
tree.subtrees.append((escape(f["name"]), f.get("size") if show_size else None))
else:
subtree, _max_string, _max_size = __construct_file_tree(
os.path.join(folder, f["name"]) if folder else f["name"],
f"[bold deep_sky_blue3]{escape(f['name'])}",
)
# Due to indentation, the filename strings of
# subdirectories are 4 characters deeper than
# their parent directories
max_string = max(max_string, _max_string + 4)
max_size = max(max_size, _max_size)
tree.subtrees.append(subtree)
return tree, max_string, max_size
def __print_project_table(self, sorted_projects, usage_info, total_size, always_show):
# Column format
column_formatting = self.__format_project_columns(
total_size=total_size, usage_info=usage_info
)
# Create table
table = Table(
title="Your Project(s)",
show_header=True,
header_style="bold",
show_footer=self.show_usage and "Usage" in column_formatting,
caption=(
"The cost is calculated from the pricing provided by Safespring (unit kr/GB/month) "
"and is therefore approximate. Contact the Data Centre for more details."
)
if self.show_usage
else None,
)
# Add columns to table
for colname, colformat in column_formatting.items():
table.add_column(
colname,
justify=colformat["justify"],
style=colformat["style"],
footer=colformat["footer"],
overflow=colformat["overflow"],
)
# Add all column values for each row to table
for proj in sorted_projects:
new_row = []
for column in column_formatting:
if column == "Size" and proj["Status"] != "Available" and not always_show:
new_row.append("---")
else:
new_row.append(
escape(
dds_cli.utils.format_api_response(
response=proj[column], key=column, binary=self.binary
)
)
)
table.add_row(*new_row)
# Print to stdout if there are any lines
dds_cli.utils.print_or_page(item=table)
def show_debug_info(self) -> None:
logging.info(
f"{EMOJI['information']} connected client version: {stylize(self.client_version, fg('green') + attr('bold'))}"
)
logging.debug("cookie: %s", self.cookie.hex())
logging.debug("kex_algorithms: %s", escape(str(self.kex_algorithms)))
logging.debug("server_host_key_algorithms: %s", self.server_host_key_algorithms)
logging.debug("encryption_algorithms_client_to_server: %s", self.encryption_algorithms_client_to_server)
logging.debug("encryption_algorithms_server_to_client: %s", self.encryption_algorithms_server_to_client)
logging.debug("mac_algorithms_client_to_server: %s", self.mac_algorithms_client_to_server)
logging.debug("mac_algorithms_server_to_client: %s", self.mac_algorithms_server_to_client)
logging.debug("compression_algorithms_client_to_server: %s", self.compression_algorithms_client_to_server)
logging.debug("compression_algorithms_server_to_client: %s", self.compression_algorithms_server_to_client)
logging.debug("languages_client_to_server: %s", self.languages_client_to_server)
logging.debug("languages_server_to_client: %s", self.languages_server_to_client)
logging.debug("first_kex_packet_follows: %s", self.first_kex_packet_follows)
def test_escape():
# Potential tags
assert escape("foo[bar]") == r"foo\[bar]"
assert escape(r"foo\[bar]") == r"foo\\\[bar]"
# Not tags (escape not required)
assert escape("[5]") == "[5]"
assert escape("\\[5]") == "\\[5]"
# Test @ escape
assert escape("[@foo]") == "\\[@foo]"
assert escape("[@]") == "\\[@]"
# https://github.com/Textualize/rich/issues/2187
assert escape("[nil, [nil]]") == r"[nil, \[nil]]"
def __post_init__(self):
self.name = person.full_name()
while "'" in self.name:
self.name = person.full_name()
self.givenName, self.surname = self.name.split()
self.userPrincipalName = f"{self.givenName}.{self.surname}@{self.domain}"
self.passwordProfile = markup.escape(person.password(length=20))
self.department = random.choice(settings.AZD_GROUP_NAMES)
self.streetAddress = f"{address.street_number()} {address.street_name()} {address.street_suffix()}"
self.state = address.state(True)
self.city = address.city()
self.postalCode = address.postal_code()
self.usageLocation = address.country()
self.country = address.country()
self.telephoneNumber = person.telephone("###-###-####")
self.mobile = person.telephone("###-###-####")
self.blockSignIn = "Yes" if self.blockSignIn else "No"
def _print(self,
val,
prefix=None,
style=None,
print_layout=True,
dont_escape=False):
if val == "--Return--":
return
if isinstance(val, str) and dont_escape is False:
val = markup.escape(val)
kwargs = {"style": str(style)} if style else {}
args = (prefix, val) if prefix else (val, )
if (show_layouts and print_layout
and self.lastcmd not in WITHOUT_LAYOUT_COMMANDS):
self._print_layout(*args, **kwargs)
else:
self.console.print(*args, **kwargs)
def list_assets_cli(models, required_models):
"""
List necessary assets.
List the assets necessary to run a given set of models.
"""
service = _configure_from_cli_arguments(models, required_models,
{"lazy_loading": True})
console = Console()
if service.configuration:
for m in service.required_models:
assets_specs = list_assets(configuration=service.configuration,
required_models=[m])
model_tree = Tree(f"[bold]{m}[/bold] ({len(assets_specs)} assets)")
if assets_specs:
for asset_spec_string in assets_specs:
model_tree.add(escape(asset_spec_string), style="dim")
console.print(model_tree)
def show_facts(self, typ: str, provider: str, long: bool):
""" Display known facts matching the criteria """
data: Dict[str, Dict[str, List[pwncat.db.Fact]]] = {}
types = typ if isinstance(typ, list) else [typ]
with Progress(
"enumerating facts",
"•",
"[cyan]{task.fields[status]}",
transient=True,
console=console,
) as progress:
task = progress.add_task("", status="initializing")
for typ in types:
for fact in pwncat.victim.enumerate.iter(
typ,
filter=lambda f: provider is None or f.source ==
provider):
progress.update(task, status=str(fact.data))
if fact.type not in data:
data[fact.type] = {}
if fact.source not in data[fact.type]:
data[fact.type][fact.source] = []
data[fact.type][fact.source].append(fact)
for typ, sources in data.items():
for source, facts in sources.items():
console.print(
f"[bright_yellow]{typ.upper()}[/bright_yellow] Facts by [blue]{source}[/blue]"
)
for fact in facts:
console.print(f" {fact.data}")
if long and getattr(fact.data, "description",
None) is not None:
console.print(
markup.escape(
textwrap.indent(fact.data.description,
" ")))
def query_yes_no(question, default="yes"):
"""Queries user for confimration"""
valid = {"yes": True, "y": True, "ye": True, "no": False, "n": False}
if default is None:
prompt = " [y/n] "
elif default == "yes":
prompt = " [Y/n] "
elif default == "no":
prompt = " [y/N] "
else:
raise ValueError("invalid default answer: '%s'" % default)
while True:
console.print(question + escape(prompt))
choice = input().lower()
if default is not None and choice == '':
return valid[default]
elif choice in valid:
return valid[choice]
else:
console.print("Please respond with 'yes' or 'no' "
"(or 'y' or 'n').\n")
def protect_and_upload(self, file, progress):
"""Process and upload the file while handling the progress bars."""
# Variables
all_ok, saved, message = (False, False, "") # Error catching
file_info = self.filehandler.data[file] # Info on current file
file_public_key, salt = ("", "") # Crypto info
# Progress bar for processing
task = progress.add_task(
description=txt.TextHandler.task_name(file=escape(file),
step="encrypt"),
total=file_info["size_raw"],
visible=not self.silent,
)
# Stream chunks from file
streamed_chunks = self.filehandler.stream_from_file(file=file)
# Stream the chunks into the encryptor to save the encrypted chunks
with fe.Encryptor(project_keys=self.keys) as encryptor:
# Encrypt and save chunks
saved, message = encryptor.encrypt_filechunks(
chunks=streamed_chunks,
outfile=file_info["path_processed"],
progress=(progress, task),
)
# Get hex version of public key -- saved in db
file_public_key = encryptor.get_public_component_hex(
private_key=encryptor.my_private)
salt = encryptor.salt
LOG.debug(
f"Updating file processed size: {file_info['path_processed']}")
# Update file info incl size, public key, salt
self.filehandler.data[file]["public_key"] = file_public_key
self.filehandler.data[file]["salt"] = salt
self.filehandler.data[file]["size_processed"] = file_info[
"path_processed"].stat().st_size
if saved:
LOG.debug(
f"File successfully encrypted: {escape(file)}. New location: {escape(str(file_info['path_processed']))}"
)
# Update progress bar for upload
progress.reset(
task,
description=txt.TextHandler.task_name(file=escape(file),
step="put"),
total=self.filehandler.data[file]["size_processed"],
step="put",
)
# Perform upload
file_uploaded, message = self.put(file=file,
progress=progress,
task=task)
# Perform db update
if file_uploaded:
db_updated, message = self.add_file_db(file=file)
if db_updated:
all_ok = True
LOG.debug(
f"File successfully uploaded and added to the database: {escape(file)}"
)
if not saved or all_ok:
# Delete temporary processed file locally
LOG.debug(
f"Deleting file {escape(str(file_info['path_processed']))} - "
f"exists: {file_info['path_processed'].exists()}")
dr.DataRemover.delete_tempfile(file=file_info["path_processed"])
# Remove progress bar task
progress.remove_task(task)
return all_ok, message
def pretty_print_type(typ):
typ_str = str(typ)
if "'" in typ_str: # for class wrapper
typ_str = typ_str.split("'")[1]
return escape(typ_str)
def escape(output: Output) -> str:
import rich.markup as markup
return markup.escape(output)
def print_escaped(self, *args, **kwargs):
self.print(escape(*args, **kwargs))
def test_escape():
assert escape("foo[bar]") == "foo[[bar]]"
def download_and_verify(self, file, progress):
"""Download the file, reveals the original data and verifies the integrity."""
all_ok, message = (False, "")
file_info = self.filehandler.data[file]
# File task for downloading
task = progress.add_task(
description=txt.TextHandler.task_name(file=escape(str(file)), step="get"),
total=file_info["size_stored"],
visible=not self.silent,
)
# Perform download
file_downloaded, message = self.get(file=file, progress=progress, task=task)
# Update progress task for decryption
progress.reset(
task,
description=txt.TextHandler.task_name(file=escape(str(file)), step="decrypt"),
total=file_info["size_original"],
)
LOG.debug(f"File {escape(str(file))} downloaded: {file_downloaded}")
if file_downloaded:
db_updated, message = self.update_db(file=file)
LOG.debug(f"Database updated: {db_updated}")
LOG.debug(f"Beginning decryption of file {escape(str(file))}...")
file_saved = False
with fe.Decryptor(
project_keys=self.keys,
peer_public=file_info["public_key"],
key_salt=file_info["salt"],
) as decryptor:
streamed_chunks = decryptor.decrypt_file(infile=file_info["path_downloaded"])
stream_to_file_func = (
fc.Compressor.decompress_filechunks
if file_info["compressed"]
else self.filehandler.write_file
)
file_saved, message = stream_to_file_func(
chunks=streamed_chunks,
outfile=file,
)
LOG.debug(f"file saved? {file_saved}")
if file_saved:
# TODO (ina): decide on checksum verification method --
# this checks original, the other is generated from compressed
all_ok, message = (
fe.Encryptor.verify_checksum(file=file, correct_checksum=file_info["checksum"])
if self.verify_checksum
else (True, "")
)
dr.DataRemover.delete_tempfile(file=file_info["path_downloaded"])
progress.remove_task(task)
return all_ok, message
def list_files(self, folder: str = None, show_size: bool = False):
"""Create a tree displaying the files within the project."""
LOG.info(f"Listing files for project '{self.project}'")
if folder:
LOG.info(f"Showing files in folder '{escape(folder)}'")
if folder is None:
folder = ""
# Make call to API
try:
response = requests.get(
DDSEndpoint.LIST_FILES,
params={"project": self.project},
json={"subpath": folder, "show_size": show_size},
headers=self.token,
timeout=DDSEndpoint.TIMEOUT,
)
except requests.exceptions.RequestException as err:
raise exceptions.APIError(
message=(
f"Failed to get list of files in project '{self.project}'"
+ (
": The database seems to be down."
if isinstance(err, requests.exceptions.ConnectionError)
else "."
)
)
)
if not response.ok:
raise exceptions.APIError(f"Failed to get list of files: '{response.text}'")
# Get response
try:
resp_json = response.json()
except simplejson.JSONDecodeError as err:
raise exceptions.APIError(f"Could not decode JSON response: '{err}'")
# Check if project empty
if "num_items" in resp_json and resp_json["num_items"] == 0:
raise exceptions.NoDataError(f"Project '{self.project}' is empty.")
# Get files
files_folders = resp_json["files_folders"]
# Sort the file/folders according to names
sorted_files_folders = sorted(files_folders, key=lambda f: f["name"])
# Create tree
tree_title = escape(folder) or f"Files / directories in project: [green]{self.project}"
tree = Tree(f"[bold magenta]{tree_title}")
if not sorted_files_folders:
raise exceptions.NoDataError(f"Could not find folder: '{escape(folder)}'")
# Get max length of file name
max_string = max([len(x["name"]) for x in sorted_files_folders])
# Get max length of size string
max_size = max(
[
len(
dds_cli.utils.format_api_response(
response=x["size"], key="Size", binary=self.binary
).split(" ", maxsplit=1)[0]
)
for x in sorted_files_folders
if show_size and "size" in x
],
default=0,
)
# Visible folders
visible_folders = []
# Add items to tree
for x in sorted_files_folders:
# Check if string is folder
is_folder = x.pop("folder")
# Att 1 for folders due to trailing /
tab = th.TextHandler.format_tabs(
string_len=len(x["name"]) + (1 if is_folder else 0),
max_string_len=max_string,
)
# Add formatting if folder and set string name
line = ""
if is_folder:
line = "[bold deep_sky_blue3]"
visible_folders.append(x["name"])
line += escape(x["name"]) + ("/" if is_folder else "")
# Add size to line if option specified
if show_size and "size" in x:
size = dds_cli.utils.format_api_response(
response=x["size"], key="Size", binary=self.binary
)
line += f"{tab}{size.split()[0]}"
# Define space between number and size format
tabs_bf_format = th.TextHandler.format_tabs(
string_len=len(size), max_string_len=max_size, tab_len=2
)
line += f"{tabs_bf_format}{size.split()[1]}"
tree.add(line)
# Print output to stdout
if len(files_folders) + 5 > dds_cli.utils.console.height:
with dds_cli.utils.console.pager():
dds_cli.utils.console.print(Padding(tree, 1))
else:
dds_cli.utils.console.print(Padding(tree, 1))
# Return variable
return visible_folders