Skip to content

Commit

Permalink
cli: refactor attach to the new approach
Browse files Browse the repository at this point in the history 
Signed-off-by: Renan Rodrigo <[email protected]>
  • Loading branch information
@renanrodrigo
renanrodrigo committed Jul 3, 2024
1 parent f019cb3 commit 84a68ca
Show file tree
Hide file tree
Showing 8 changed files with 300 additions and 344 deletions.
23 changes: 23 additions & 0 deletions features/help.feature
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,29 @@ Feature: Pro Client help text
--variant VARIANT The name of the variant to use when enabling the
service
"""
When I run `pro attach --help` as non-root
Then stdout matches regexp:
"""
usage: pro attach \[flags\]
Attach this machine to an Ubuntu Pro subscription with a token obtained from:
https://ubuntu.com/pro/dashboard
When running this command without a token, it will generate a short code
and prompt you to attach the machine to your Ubuntu Pro account using
a web browser.
positional arguments:
token token obtained for Ubuntu Pro authentication
(optional arguments|options):
-h, --help show this help message and exit
--no-auto-enable do not enable any recommended services automatically
--attach-config ATTACH_CONFIG
use the provided attach config file instead of passing
the token on the cli
--format \{cli,json\} output in the specified format \(default: cli\)
"""

Examples: ubuntu release
| release | machine_type |
Expand Down
194 changes: 12 additions & 182 deletions uaclient/cli/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
lock,
log,
messages,
secret_manager,
security_status,
status,
timer,
Expand All @@ -33,39 +32,23 @@
FullAutoAttachOptions,
_full_auto_attach,
)
from uaclient.api.u.pro.attach.magic.initiate.v1 import _initiate
from uaclient.api.u.pro.attach.magic.revoke.v1 import (
MagicAttachRevokeOptions,
_revoke,
)
from uaclient.api.u.pro.attach.magic.wait.v1 import (
MagicAttachWaitOptions,
_wait,
)
from uaclient.api.u.pro.security.status.reboot_required.v1 import (
_reboot_required,
)
from uaclient.apt import AptProxyScope, setup_apt_proxy
from uaclient.cli import cli_util, fix
from uaclient.cli.api import api_command
from uaclient.cli.attach import attach_command
from uaclient.cli.collect_logs import collect_logs_command
from uaclient.cli.constants import NAME, USAGE_TMPL
from uaclient.cli.disable import disable_command, perform_disable
from uaclient.cli.enable import enable_command
from uaclient.data_types import AttachActionsConfigFile, IncorrectTypeError
from uaclient.entitlements import (
create_enable_entitlements_not_found_error,
entitlements_disable_order,
get_valid_entitlement_names,
)
from uaclient.entitlements.entitlement_status import (
ApplicationStatus,
CanEnableFailure,
)
from uaclient.entitlements import entitlements_disable_order
from uaclient.entitlements.entitlement_status import ApplicationStatus
from uaclient.files import machine_token, state_files
from uaclient.log import get_user_or_root_log_file_path
from uaclient.timer.update_messaging import refresh_motd, update_motd_messages
from uaclient.yaml import safe_dump, safe_load
from uaclient.yaml import safe_dump

UA_AUTH_TOKEN_URL = "https://auth.contracts.canonical.com"

Expand All @@ -74,7 +57,13 @@
event = event_logger.get_event_logger()
LOG = logging.getLogger(util.replace_top_level_logger_name(__name__))

COMMANDS = [api_command, collect_logs_command, disable_command, enable_command]
COMMANDS = [
api_command,
attach_command,
collect_logs_command,
disable_command,
enable_command,
]


class UAArgumentParser(argparse.ArgumentParser):
Expand Down Expand Up @@ -190,35 +179,6 @@ def config_parser(parser):
return parser


def attach_parser(parser):
"""Build or extend an arg parser for attach subcommand."""
parser.usage = USAGE_TMPL.format(name=NAME, command="attach <token>")
parser.formatter_class = argparse.RawDescriptionHelpFormatter
parser.prog = "attach"
parser.description = messages.CLI_ATTACH_DESC
parser._optionals.title = messages.CLI_FLAGS
parser.add_argument("token", nargs="?", help=messages.CLI_ATTACH_TOKEN)
parser.add_argument(
"--no-auto-enable",
action="store_false",
dest="auto_enable",
help=messages.CLI_ATTACH_NO_AUTO_ENABLE,
)
parser.add_argument(
"--attach-config",
type=argparse.FileType("r"),
help=messages.CLI_ATTACH_ATTACH_CONFIG,
)
parser.add_argument(
"--format",
action="store",
choices=["cli", "json"],
default="cli",
help=messages.CLI_FORMAT_DESC.format(default="cli"),
)
return parser


def security_status_parser(parser):
"""Build or extend an arg parser for security-status subcommand."""
parser.prog = "security-status"
Expand Down Expand Up @@ -719,26 +679,6 @@ def _detach(cfg: config.UAConfig, assume_yes: bool, json_output: bool) -> int:
return 0


def _post_cli_attach(cfg: config.UAConfig) -> None:
machine_token_file = machine_token.get_machine_token_file(cfg)
contract_name = machine_token_file.contract_name

if contract_name:
event.info(
messages.ATTACH_SUCCESS_TMPL.format(contract_name=contract_name)
)
else:
event.info(messages.ATTACH_SUCCESS_NO_CONTRACT_NAME)

daemon.stop()
daemon.cleanup(cfg)

status_dict, _ret = actions.status(cfg)
output = status.format_tabular(status_dict)
event.info(util.handle_unicode_characters(output))
event.process_events()


@cli_util.assert_root
def action_auto_attach(args, *, cfg: config.UAConfig, **kwargs) -> int:
try:
Expand All @@ -751,114 +691,10 @@ def action_auto_attach(args, *, cfg: config.UAConfig, **kwargs) -> int:
event.info(messages.E_ATTACH_FAILURE.msg)
return 1
else:
_post_cli_attach(cfg)
cli_util.post_cli_attach(cfg)
return 0


def _magic_attach(args, *, cfg, **kwargs):
if args.format == "json":
raise exceptions.MagicAttachInvalidParam(
param="--format",
value=args.format,
)

event.info(messages.CLI_MAGIC_ATTACH_INIT)
initiate_resp = _initiate(cfg=cfg)
event.info(
"\n"
+ messages.CLI_MAGIC_ATTACH_SIGN_IN.format(
user_code=initiate_resp.user_code
)
)

wait_options = MagicAttachWaitOptions(magic_token=initiate_resp.token)

try:
wait_resp = _wait(options=wait_options, cfg=cfg)
except exceptions.MagicAttachTokenError as e:
event.info(messages.CLI_MAGIC_ATTACH_FAILED)

revoke_options = MagicAttachRevokeOptions(
magic_token=initiate_resp.token
)
_revoke(options=revoke_options, cfg=cfg)
raise e

event.info("\n" + messages.CLI_MAGIC_ATTACH_PROCESSING)
return wait_resp.contract_token


@cli_util.assert_not_attached
@cli_util.assert_root
@cli_util.assert_lock_file("pro attach")
def action_attach(args, *, cfg, **kwargs):
if args.token and args.attach_config:
raise exceptions.CLIAttachTokenArgXORConfig()
elif not args.token and not args.attach_config:
token = _magic_attach(args, cfg=cfg)
enable_services_override = None
elif args.token:
token = args.token
secret_manager.secrets.add_secret(token)
enable_services_override = None
else:
try:
attach_config = AttachActionsConfigFile.from_dict(
safe_load(args.attach_config)
)
except IncorrectTypeError as e:
raise exceptions.AttachInvalidConfigFileError(
config_name=args.attach_config.name, error=e.msg
)

token = attach_config.token
enable_services_override = attach_config.enable_services

allow_enable = args.auto_enable and enable_services_override is None

try:
actions.attach_with_token(cfg, token=token, allow_enable=allow_enable)
except exceptions.ConnectivityError:
raise exceptions.AttachError()
else:
ret = 0
if enable_services_override is not None and args.auto_enable:
found, not_found = get_valid_entitlement_names(
enable_services_override, cfg
)
for name in found:
ent_ret, reason = actions.enable_entitlement_by_name(cfg, name)
if not ent_ret:
ret = 1
if (
reason is not None
and isinstance(reason, CanEnableFailure)
and reason.message is not None
):
event.info(reason.message.msg)
event.error(
error_msg=reason.message.msg,
error_code=reason.message.name,
service=name,
)
else:
event.service_processed(name)

if not_found:
error = create_enable_entitlements_not_found_error(
not_found, cfg=cfg
)
event.info(error.msg, file_type=sys.stderr)
event.error(error_msg=error.msg, error_code=error.msg_code)
ret = 1

contract_client = contract.UAContractClient(cfg)
contract_client.update_activity_token()

_post_cli_attach(cfg)
return ret


def get_parser(cfg: config.UAConfig):
parser = UAArgumentParser(
prog=NAME,
Expand All @@ -885,12 +721,6 @@ def get_parser(cfg: config.UAConfig):
for command in COMMANDS:
command.register(subparsers)

parser_attach = subparsers.add_parser(
"attach", help=messages.CLI_ROOT_ATTACH
)
attach_parser(parser_attach)
parser_attach.set_defaults(action=action_attach)

parser_auto_attach = subparsers.add_parser(
"auto-attach", help=messages.CLI_ROOT_AUTO_ATTACH
)
Expand Down
Loading

0 comments on commit 84a68ca

Please sign in to comment.