Files
onefuzz/src/deployment/deploylib/registration.py
2023-02-01 20:49:14 +00:00

917 lines
29 KiB
Python

#!/usr/bin/env python
#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import argparse
import logging
import re
import time
import urllib.parse
from datetime import datetime, timedelta
from enum import Enum
from typing import (
Any,
Callable,
Dict,
List,
NamedTuple,
Optional,
Tuple,
TypeVar,
Union,
)
from uuid import UUID
import requests
from azure.common.credentials import get_cli_profile
from functional import seq
from msrest.serialization import TZ_UTC
logger = logging.getLogger("deploy")
## https://docs.microsoft.com/en-us/graph/api/overview?view=graph-rest-1.0
GRAPH_RESOURCE = "https://graph.microsoft.com"
GRAPH_RESOURCE_ENDPOINT = "https://graph.microsoft.com/v1.0"
NameOrAppId = Union[str, UUID]
class GraphQueryError(Exception):
def __init__(self, message: str, status_code: Optional[int]) -> None:
super(GraphQueryError, self).__init__(message)
self.message = message
self.status_code = status_code
## Queries microsoft graph api and return
def query_microsoft_graph(
method: str,
resource: str,
params: Optional[Dict] = None,
body: Optional[Dict] = None,
subscription: Optional[str] = None,
) -> Dict:
profile = get_cli_profile()
(token_type, access_token, _), _, _ = profile.get_raw_token(
resource=GRAPH_RESOURCE, subscription=subscription
)
url = urllib.parse.urljoin(f"{GRAPH_RESOURCE_ENDPOINT}/", resource)
headers = {
"Authorization": "%s %s" % (token_type, access_token),
"Content-Type": "application/json",
}
response = requests.request(
method=method, url=url, headers=headers, params=params, json=body
)
if 200 <= response.status_code < 300:
if response.content and response.content.strip():
json = response.json()
if isinstance(json, Dict):
return json
else:
raise GraphQueryError(
f"invalid data received expected a json object: HTTP {response.status_code} - {json}",
response.status_code,
)
else:
return {}
else:
error_text = str(response.content, encoding="utf-8", errors="backslashreplace")
raise GraphQueryError(
f"request did not succeed: HTTP {response.status_code} - {error_text}",
response.status_code,
)
def query_microsoft_graph_list(
method: str,
resource: str,
params: Optional[Dict] = None,
body: Optional[Dict] = None,
subscription: Optional[str] = None,
) -> List[Any]:
result = query_microsoft_graph(
method,
resource,
params,
body,
subscription,
)
value = result.get("value")
if isinstance(value, list):
return value
else:
raise GraphQueryError("Expected data containing a list of values", None)
def get_tenant_id(subscription_id: Optional[str] = None) -> str:
profile = get_cli_profile()
_, _, tenant_id = profile.get_raw_token(
resource=GRAPH_RESOURCE, subscription=subscription_id
)
if isinstance(tenant_id, str):
return tenant_id
else:
raise Exception(
f"unable to retrive tenant_id for subscription {subscription_id}"
)
OperationResult = TypeVar("OperationResult")
def retry(
operation: Callable[[Any], OperationResult],
description: str,
tries: int = 10,
wait_duration: int = 10,
data: Any = None,
) -> OperationResult:
count = 0
error = None
while True:
try:
return operation(data)
except GraphQueryError as err:
error = err
# modeled after AZ-CLI's handling of missing application
# See: https://github.com/Azure/azure-cli/blob/
# e015d5bcba0c2d21dc42189daa43dc1eb82d2485/src/azure-cli/
# azure/cli/command_modules/util/tests/
# latest/test_rest.py#L191-L192
if "Request_ResourceNotFound" in repr(err):
logger.info(f"failed '{description}' missing required resource")
else:
logger.warning(f"failed '{description}': {err.message}")
except Exception as exc:
exception = exc
logger.error(f"failed '{description}'. logging stack trace.")
logger.error(exc)
count += 1
if count >= tries:
if error:
raise error
elif exception:
raise exception
else:
raise Exception(f"failed '{description}'")
else:
logger.info(
f"waiting {wait_duration} seconds before retrying '{description}'"
)
time.sleep(wait_duration)
class ApplicationInfo(NamedTuple):
client_id: UUID
client_secret: str
authority: str
class OnefuzzAppRole(Enum):
ManagedNode = "ManagedNode"
CliClient = "CliClient"
UserAssignment = "UserAssignment"
UnmanagedNode = "UnmanagedNode"
def register_application(
registration_name: str,
onefuzz_instance_name: str,
approle: OnefuzzAppRole,
subscription_id: str,
) -> ApplicationInfo:
logger.info("retrieving the application registration %s" % registration_name)
app = get_application(
display_name=registration_name, subscription_id=subscription_id
)
if not app:
logger.info("No existing registration found. creating a new one")
app = create_application_registration(
onefuzz_instance_name, registration_name, approle, subscription_id
)
else:
logger.info(
"Found existing application objectId '%s' - appid '%s'"
% (app["id"], app["appId"])
)
onefuzz_app = get_application(
display_name=onefuzz_instance_name, subscription_id=subscription_id
)
if not (onefuzz_app):
raise Exception("onefuzz app not found")
pre_authorized_applications = onefuzz_app["api"]["preAuthorizedApplications"]
if app["appId"] not in [app["appId"] for app in pre_authorized_applications]:
authorize_application(UUID(app["appId"]), UUID(onefuzz_app["appId"]))
password = create_application_credential(registration_name, subscription_id)
tenant_id = get_tenant_id(subscription_id=subscription_id)
return ApplicationInfo(
client_id=app["appId"],
client_secret=password,
authority=("https://login.microsoftonline.com/%s" % tenant_id),
)
def create_application_credential(application_name: str, subscription_id: str) -> str:
"""Add a new password to the application registration"""
logger.info("creating application credential for '%s'" % application_name)
app = get_application(display_name=application_name)
if not app:
raise Exception("app not found")
(_, password) = add_application_password(
f"{application_name}_password", app["id"], subscription_id
)
return str(password)
def create_application_registration(
onefuzz_instance_name: str, name: str, approle: OnefuzzAppRole, subscription_id: str
) -> Any:
"""Create an application registration"""
app = get_application(
display_name=onefuzz_instance_name, subscription_id=subscription_id
)
if not app:
raise Exception("onefuzz app registration not found")
resource_access = [
{"id": role["id"], "type": "Scope"}
for role in app["appRoles"]
if role["value"] == approle.value
]
params = {
"isDeviceOnlyAuthSupported": True,
"displayName": name,
"isFallbackPublicClient": True,
"requiredResourceAccess": (
[
{
"resourceAccess": resource_access,
"resourceAppId": app["appId"],
}
]
if len(resource_access) > 0
else []
),
}
registered_app = query_microsoft_graph(
method="POST",
resource="applications",
body=params,
subscription=subscription_id,
)
# next patch the redirect URIs; we must do this
# separately because we need the AppID to include
query_microsoft_graph(
method="PATCH",
resource=f"applications/{registered_app['id']}",
body={
"publicClient": {
"redirectUris": [
"https://%s.azurewebsites.net" % onefuzz_instance_name,
"http://localhost", # required for browser auth
f"ms-appx-web://Microsoft.AAD.BrokerPlugin/{app['appId']}", # required for broker auth
]
},
},
subscription=subscription_id,
)
logger.info("creating service principal")
service_principal_params = {
"accountEnabled": True,
"appRoleAssignmentRequired": False,
"servicePrincipalType": "Application",
"appId": registered_app["appId"],
}
def try_sp_create() -> None:
error: Optional[Exception] = None
for _ in range(10):
try:
query_microsoft_graph(
method="POST",
resource="servicePrincipals",
body=service_principal_params,
subscription=subscription_id,
)
return
except GraphQueryError as err:
# work around timing issue when creating service principal
# https://github.com/Azure/azure-cli/issues/14767
if (
"service principal being created must in the local tenant"
not in str(err)
):
raise err
logger.warning(
"creating service principal failed with an error that occurs "
"due to AAD race conditions"
)
time.sleep(60)
if error is None:
raise Exception("service principal creation failed")
else:
raise error
try_sp_create()
registered_app_id = registered_app["appId"]
app_id = app["appId"]
authorize_and_assign_role(app_id, registered_app_id, approle, subscription_id)
return registered_app
def authorize_and_assign_role(
onfuzz_app_id: UUID,
registered_app_id: UUID,
role: OnefuzzAppRole,
subscription_id: str,
) -> None:
def try_authorize_application(data: Any) -> None:
authorize_application(
registered_app_id,
onfuzz_app_id,
subscription_id=subscription_id,
)
retry(try_authorize_application, "authorize application")
def try_assign_instance_role(data: Any) -> None:
assign_instance_app_role(
onfuzz_app_id, registered_app_id, subscription_id, role
)
retry(try_assign_instance_role, "assingn role")
def add_application_password(
password_name: str, app_object_id: UUID, subscription_id: str
) -> Tuple[str, str]:
def create_password(data: Any) -> Tuple[str, str]:
password = add_application_password_impl(
password_name, app_object_id, subscription_id
)
logger.info("app password created")
return password
# Work-around the race condition where the app is created but passwords cannot
# be created yet.
return retry(create_password, "create password")
def add_application_password_impl(
password_name: str, app_object_id: UUID, subscription_id: str
) -> Tuple[str, str]:
app = query_microsoft_graph(
method="GET",
resource="applications/%s" % app_object_id,
subscription=subscription_id,
)
passwords = [
x for x in app["passwordCredentials"] if x["displayName"] == password_name
]
if len(passwords) > 0:
key_id = passwords[0]["keyId"]
query_microsoft_graph(
method="POST",
resource="applications/%s/removePassword" % app_object_id,
body={"keyId": key_id},
subscription=subscription_id,
)
password_request = {
"passwordCredential": {
"displayName": "%s" % password_name,
"startDateTime": "%s" % datetime.now(TZ_UTC).strftime("%Y-%m-%dT%H:%M:%SZ"),
"endDateTime": "%s"
% (datetime.now(TZ_UTC) + timedelta(days=365)).strftime(
"%Y-%m-%dT%H:%M:%SZ"
),
}
}
password: Dict = query_microsoft_graph(
method="POST",
resource="applications/%s/addPassword" % app_object_id,
body=password_request,
subscription=subscription_id,
)
return (password_name, password["secretText"])
def get_application(
app_id: Optional[UUID] = None,
display_name: Optional[str] = None,
subscription_id: Optional[str] = None,
) -> Optional[Any]:
filters = []
if app_id:
filters.append("appId eq '%s'" % app_id)
if display_name:
filters.append("displayName eq '%s'" % display_name)
filter_str = " and ".join(filters)
apps = query_microsoft_graph(
method="GET",
resource="applications",
params={
"$filter": filter_str,
},
subscription=subscription_id,
)
number_of_apps = len(apps["value"])
if number_of_apps == 0:
return None
elif number_of_apps == 1:
return apps["value"][0]
else:
raise Exception(
f"Found {number_of_apps} application matching filter: '{filter_str}'"
)
def authorize_application(
registration_app_id: UUID,
onefuzz_app_id: UUID,
permissions: List[str] = ["user_impersonation"],
subscription_id: Optional[str] = None,
) -> None:
logger.info(
f"authorizing registration {registration_app_id} to access application {onefuzz_app_id} with the permissions '{', '.join(permissions)}'"
)
onefuzz_app = get_application(
app_id=onefuzz_app_id, subscription_id=subscription_id
)
if onefuzz_app is None:
logger.error("Application '%s' not found", onefuzz_app_id)
return
scopes = seq(onefuzz_app["api"]["oauth2PermissionScopes"]).filter(
lambda scope: scope["value"] in permissions
)
existing_preAuthorizedApplications = (
seq(onefuzz_app["api"]["preAuthorizedApplications"])
.map(
lambda paa: seq(paa["delegatedPermissionIds"]).map(
lambda permission_id: (paa["appId"], permission_id)
)
)
.flatten()
)
preAuthorizedApplications = (
scopes.map(lambda s: (str(registration_app_id), s["id"]))
.union(existing_preAuthorizedApplications)
.distinct()
.group_by_key()
.map(lambda data: {"appId": data[0], "delegatedPermissionIds": data[1]})
)
onefuzz_app_id = onefuzz_app["id"]
def add_preauthorized_app(app_list: List[Dict]) -> None:
try:
query_microsoft_graph(
method="PATCH",
resource="applications/%s" % onefuzz_app_id,
body={"api": {"preAuthorizedApplications": app_list}},
subscription=subscription_id,
)
except GraphQueryError as e:
m = re.search(
"Property PreAuthorizedApplication references "
"applications (.*?) that cannot be found.",
e.message,
)
if m:
invalid_app_id = m.group(1)
if invalid_app_id:
for app in app_list:
if app["appId"] == invalid_app_id:
logger.warning(
f"removing invalid id {invalid_app_id} "
"for the next request"
)
app_list.remove(app)
raise e
retry(
add_preauthorized_app,
"authorize application",
data=preAuthorizedApplications.to_list(),
)
def create_and_display_registration(
onefuzz_instance_name: str,
registration_name: str,
approle: OnefuzzAppRole,
subscription_id: str,
*,
display_secret: bool = False,
) -> None:
logger.info("Updating application registration")
application_info = register_application(
registration_name=registration_name,
onefuzz_instance_name=onefuzz_instance_name,
approle=approle,
subscription_id=subscription_id,
)
if display_secret:
print("Registration complete")
print("These generated credentials are valid for a year")
print(f"client_id: {application_info.client_id}")
print(f"client_secret: {application_info.client_secret}")
def update_pool_registration(onefuzz_instance_name: str, subscription_id: str) -> None:
create_and_display_registration(
onefuzz_instance_name,
"%s_pool" % onefuzz_instance_name,
OnefuzzAppRole.ManagedNode,
subscription_id,
)
def assign_app_role(
principal_id: str,
application_id: str,
role_names: List[str],
subscription_id: str,
) -> None:
application_registrations = query_microsoft_graph_list(
method="GET",
resource="servicePrincipals",
params={
"$filter": f"appId eq '{application_id}'",
},
subscription=subscription_id,
)
if len(application_registrations) == 0:
raise Exception(f"appid '{application_id}' was not found:")
app = application_registrations[0]
roles = (
seq(app["appRoles"]).filter(lambda role: role["value"] in role_names).to_list()
)
if len(roles) < len(role_names):
existing_roles = [role["value"] for role in roles]
missing_roles = [
role_name for role_name in role_names if role_name not in existing_roles
]
raise Exception(
f"The following roles could not be found in appId '{application_id}': {missing_roles}"
)
expected_role_ids = [role["id"] for role in roles]
assignments = query_microsoft_graph_list(
method="GET",
resource=f"servicePrincipals/{principal_id}/appRoleAssignments",
subscription=subscription_id,
)
assigned_role_ids = [assignment["appRoleId"] for assignment in assignments]
missing_assignments = [
id for id in expected_role_ids if id not in assigned_role_ids
]
if missing_assignments:
for app_role_id in missing_assignments:
query_microsoft_graph(
method="POST",
resource=f"servicePrincipals/{principal_id}/appRoleAssignedTo",
body={
"principalId": principal_id,
"resourceId": app["id"],
"appRoleId": app_role_id,
},
subscription=subscription_id,
)
def assign_instance_app_role(
onefuzz_instance: NameOrAppId,
application_name: NameOrAppId,
subscription_id: str,
app_role: OnefuzzAppRole,
) -> None:
"""
Allows the application to access the service by assigning
their managed identity to the provided App Role
"""
logger.info(
f"Assigning app role {app_role} from {onefuzz_instance} to {application_name}"
)
if isinstance(onefuzz_instance, str):
onefuzz_service_appIds = query_microsoft_graph_list(
method="GET",
resource="applications",
params={
"$filter": "displayName eq '%s'" % onefuzz_instance,
"$select": "appId",
},
subscription=subscription_id,
)
if len(onefuzz_service_appIds) == 0:
raise Exception("onefuzz app registration not found")
appId = onefuzz_service_appIds[0]["appId"]
else:
appId = onefuzz_instance
onefuzz_service_principals = query_microsoft_graph_list(
method="GET",
resource="servicePrincipals",
params={"$filter": "appId eq '%s'" % appId},
subscription=subscription_id,
)
if len(onefuzz_service_principals) == 0:
raise Exception("onefuzz app service principal not found")
onefuzz_service_principal = onefuzz_service_principals[0]
if isinstance(application_name, str):
application_service_principals = query_microsoft_graph_list(
method="GET",
resource="servicePrincipals",
params={"$filter": "displayName eq '%s'" % application_name},
subscription=subscription_id,
)
else:
application_service_principals = query_microsoft_graph_list(
method="GET",
resource="servicePrincipals",
params={"$filter": "appId eq '%s'" % application_name},
subscription=subscription_id,
)
if len(application_service_principals) == 0:
raise Exception(f"application '{application_name}' service principal not found")
application_service_principal = application_service_principals[0]
managed_node_role = (
seq(onefuzz_service_principal["appRoles"])
.filter(lambda x: x["value"] == app_role.value)
.head_option()
)
if not managed_node_role:
raise Exception(
f"{app_role.value} role not found in the OneFuzz application "
"registration. Please redeploy the instance"
)
assignments = query_microsoft_graph_list(
method="GET",
resource="servicePrincipals/%s/appRoleAssignments"
% application_service_principal["id"],
subscription=subscription_id,
)
# check if the role is already assigned
role_assigned = seq(assignments).find(
lambda assignment: assignment["appRoleId"] == managed_node_role["id"]
)
if not role_assigned:
query_microsoft_graph(
method="POST",
resource="servicePrincipals/%s/appRoleAssignedTo"
% application_service_principal["id"],
body={
"principalId": application_service_principal["id"],
"resourceId": onefuzz_service_principal["id"],
"appRoleId": managed_node_role["id"],
},
subscription=subscription_id,
)
def set_app_audience(
objectId: str, audience: str, subscription_id: Optional[str] = None
) -> None:
# typical audience values: AzureADMyOrg, AzureADMultipleOrgs
http_body = {"signInAudience": audience}
try:
query_microsoft_graph(
method="PATCH",
resource="applications/%s" % objectId,
body=http_body,
subscription=subscription_id,
)
except GraphQueryError:
query = (
"az rest --method patch --url "
"https://graph.microsoft.com/v1.0/applications/%s "
"--body '%s' --headers \"Content-Type\"=application/json"
% (objectId, http_body)
)
logger.warning(
"execute the following query in the azure portal bash shell and "
"run deploy.py again : \n%s",
query,
)
err_str = (
"Unable to set signInAudience using Microsoft Graph Query API. \n"
"The user must enable single/multi tenancy in the "
"'Authentication' blade of the "
"Application Registration in the "
"AAD web portal, or use the azure bash shell "
"using the command given above."
)
raise Exception(err_str)
def get_signed_in_user(subscription_id: Optional[str]) -> Any:
# Get principalId by retrieving owner for SP
try:
app = query_microsoft_graph(
method="GET",
resource="me/",
subscription=subscription_id,
)
return app
except GraphQueryError:
query = (
"az rest --method post --url "
"https://graph.microsoft.com/v1.0/me "
'--headers "Content-Type"=application/json'
)
logger.warning(
"execute the following query in the azure portal bash shell and "
"run deploy.py again : \n%s",
query,
)
err_str = "Unable to retrieve signed-in user via Microsoft Graph Query API. \n"
raise Exception(err_str)
def get_service_principal(app_id: str, subscription_id: Optional[str]) -> Any:
try:
service_principals = query_microsoft_graph_list(
method="GET",
resource="servicePrincipals",
params={"$filter": f"appId eq '{app_id}'"},
subscription=subscription_id,
)
if len(service_principals) != 0:
return service_principals[0]
else:
raise GraphQueryError(
f"Could not retrieve any service principals for App Id: {app_id}", 400
)
except GraphQueryError:
err_str = "Unable to add retrieve SP using Microsoft Graph Query API. \n"
raise Exception(err_str)
def add_user(object_id: str, principal_id: str, role_id: str) -> None:
# Get principalId by retrieving owner for SP
# need to add users with proper role assignment
http_body = {
"principalId": principal_id,
"resourceId": object_id,
"appRoleId": role_id,
}
try:
query_microsoft_graph(
method="POST",
resource="users/%s/appRoleAssignments" % principal_id,
body=http_body,
)
except GraphQueryError as ex:
if "Permission being assigned already exists" not in ex.message:
query = (
"az rest --method post --url "
"https://graph.microsoft.com/v1.0/users/%s/appRoleAssignments "
"--body '%s' --headers \"Content-Type\"=application/json"
% (principal_id, http_body)
)
logger.warning(
"execute the following query in the azure portal bash shell and "
"run deploy.py again : \n%s",
query,
)
err_str = "Unable to add user to SP using Microsoft Graph Query API. \n"
raise Exception(err_str)
else:
logger.info("User already assigned to application.")
def main() -> None:
formatter = argparse.ArgumentDefaultsHelpFormatter
parent_parser = argparse.ArgumentParser(add_help=False)
parent_parser.add_argument(
"onefuzz_instance", help="the name of the onefuzz instance"
)
parent_parser.add_argument("subscription_id")
parser = argparse.ArgumentParser(
formatter_class=formatter,
description=(
"Create an application registration and/or "
"generate a password for the pool agent"
),
)
parser.add_argument("-v", "--verbose", action="store_true")
subparsers = parser.add_subparsers(title="commands", dest="command")
subparsers.add_parser("update_pool_registration", parents=[parent_parser])
scaleset_role_assignment_parser = subparsers.add_parser(
"assign_scaleset_role",
parents=[parent_parser],
)
scaleset_role_assignment_parser.add_argument(
"--scaleset_name",
help="the name of the scaleset",
)
cli_registration_parser = subparsers.add_parser(
"create_cli_registration", parents=[parent_parser]
)
cli_registration_parser.add_argument(
"--registration_name", help="the name of the cli registration"
)
register_app_parser = subparsers.add_parser("register_app", parents=[parent_parser])
register_app_parser.add_argument(
"--app_id", help="the application id to register", required=True
)
register_app_parser.add_argument(
"--role",
help=f"the role of the application to register. Valid values: {', '.join([member.value for member in OnefuzzAppRole])}",
required=True,
)
args = parser.parse_args()
if args.verbose:
level = logging.DEBUG
else:
level = logging.WARN
logging.basicConfig(format="%(levelname)s:%(message)s", level=level)
logging.getLogger("deploy").setLevel(logging.INFO)
onefuzz_instance_name = args.onefuzz_instance
if args.command == "update_pool_registration":
update_pool_registration(onefuzz_instance_name, args.subscription_id)
elif args.command == "create_cli_registration":
registration_name = args.registration_name or ("%s_cli" % onefuzz_instance_name)
create_and_display_registration(
onefuzz_instance_name,
registration_name,
OnefuzzAppRole.CliClient,
args.subscription_id,
display_secret=True,
)
elif args.command == "register_app":
onefuzz_app_id = get_application(display_name=onefuzz_instance_name)
if not onefuzz_app_id:
raise Exception("could not find onefuzz application")
authorize_and_assign_role(
UUID(onefuzz_app_id["appId"]),
UUID(args.app_id),
OnefuzzAppRole(args.role),
args.subscription_id,
)
elif args.command == "assign_scaleset_role":
assign_instance_app_role(
onefuzz_instance_name,
args.scaleset_name,
args.subscription_id,
OnefuzzAppRole.ManagedNode,
)
else:
raise Exception("invalid arguments")
if __name__ == "__main__":
main()