moto/moto/organizations/models.py

from __future__ import unicode_literals

import datetime
import re

from moto.core import BaseBackend, BaseModel
from moto.core.exceptions import RESTError
from moto.core.utils import unix_time
from moto.organizations import utils


class FakeOrganization(BaseModel):

    def __init__(self, feature_set):
        self.id = utils.make_random_org_id()
        self.root_id = utils.make_random_root_id()
        self.feature_set = feature_set
        self.master_account_id = utils.MASTER_ACCOUNT_ID
        self.master_account_email = utils.MASTER_ACCOUNT_EMAIL
        self.available_policy_types = [{
            'Type': 'SERVICE_CONTROL_POLICY',
            'Status': 'ENABLED'
        }]

    @property
    def arn(self):
        return utils.ORGANIZATION_ARN_FORMAT.format(self.master_account_id, self.id)

    @property
    def master_account_arn(self):
        return utils.MASTER_ACCOUNT_ARN_FORMAT.format(self.master_account_id, self.id)

    def describe(self):
        return {
            'Organization': {
                'Id': self.id,
                'Arn': self.arn,
                'FeatureSet': self.feature_set,
                'MasterAccountArn': self.master_account_arn,
                'MasterAccountId': self.master_account_id,
                'MasterAccountEmail': self.master_account_email,
                'AvailablePolicyTypes': self.available_policy_types,
            }
        }


class FakeAccount(BaseModel):

    def __init__(self, organization, **kwargs):
        self.type = 'ACCOUNT'
        self.organization_id = organization.id
        self.master_account_id = organization.master_account_id
        self.create_account_status_id = utils.make_random_create_account_status_id()
        self.id = utils.make_random_account_id()
        self.name = kwargs['AccountName']
        self.email = kwargs['Email']
        self.create_time = datetime.datetime.utcnow()
        self.status = 'ACTIVE'
        self.joined_method = 'CREATED'
        self.parent_id = organization.root_id
        self.attached_policies = []

    @property
    def arn(self):
        return utils.ACCOUNT_ARN_FORMAT.format(
            self.master_account_id,
            self.organization_id,
            self.id
        )

    @property
    def create_account_status(self):
        return {
            'CreateAccountStatus': {
                'Id': self.create_account_status_id,
                'AccountName': self.name,
                'State': 'SUCCEEDED',
                'RequestedTimestamp': unix_time(self.create_time),
                'CompletedTimestamp': unix_time(self.create_time),
                'AccountId': self.id,
            }
        }

    def describe(self):
        return {
            'Account': {
                'Id': self.id,
                'Arn': self.arn,
                'Email': self.email,
                'Name': self.name,
                'Status': self.status,
                'JoinedMethod': self.joined_method,
                'JoinedTimestamp': unix_time(self.create_time),
            }
        }


class FakeOrganizationalUnit(BaseModel):

    def __init__(self, organization, **kwargs):
        self.type = 'ORGANIZATIONAL_UNIT'
        self.organization_id = organization.id
        self.master_account_id = organization.master_account_id
        self.id = utils.make_random_ou_id(organization.root_id)
        self.name = kwargs.get('Name')
        self.parent_id = kwargs.get('ParentId')
        self._arn_format = utils.OU_ARN_FORMAT
        self.attached_policies = []

    @property
    def arn(self):
        return self._arn_format.format(
            self.master_account_id,
            self.organization_id,
            self.id
        )

    def describe(self):
        return {
            'OrganizationalUnit': {
                'Id': self.id,
                'Arn': self.arn,
                'Name': self.name,
            }
        }


class FakeRoot(FakeOrganizationalUnit):

    def __init__(self, organization, **kwargs):
        super(FakeRoot, self).__init__(organization, **kwargs)
        self.type = 'ROOT'
        self.id = organization.root_id
        self.name = 'Root'
        self.policy_types = [{
            'Type': 'SERVICE_CONTROL_POLICY',
            'Status': 'ENABLED'
        }]
        self._arn_format = utils.ROOT_ARN_FORMAT
        self.attached_policies = []

    def describe(self):
        return {
            'Id': self.id,
            'Arn': self.arn,
            'Name': self.name,
            'PolicyTypes': self.policy_types
        }


class FakeServiceControlPolicy(BaseModel):

    def __init__(self, organization, **kwargs):
        self.type = 'POLICY'
        self.content = kwargs.get('Content')
        self.description = kwargs.get('Description')
        self.name = kwargs.get('Name')
        self.type = kwargs.get('Type')
        self.id = utils.make_random_service_control_policy_id()
        self.aws_managed = False
        self.organization_id = organization.id
        self.master_account_id = organization.master_account_id
        self._arn_format = utils.SCP_ARN_FORMAT
        self.attachments = []

    @property
    def arn(self):
        return self._arn_format.format(
            self.master_account_id,
            self.organization_id,
            self.id
        )

    def describe(self):
        return {
            'Policy': {
                'PolicySummary': {
                    'Id': self.id,
                    'Arn': self.arn,
                    'Name': self.name,
                    'Description': self.description,
                    'Type': self.type,
                    'AwsManaged': self.aws_managed,
                },
                'Content': self.content
            }
        }


class OrganizationsBackend(BaseBackend):

    def __init__(self):
        self.org = None
        self.accounts = []
        self.ou = []
        self.policies = []

    def create_organization(self, **kwargs):
        self.org = FakeOrganization(kwargs['FeatureSet'])
        self.ou.append(FakeRoot(self.org))
        return self.org.describe()

    def describe_organization(self):
        if not self.org:
            raise RESTError(
                'AWSOrganizationsNotInUseException',
                "Your account is not a member of an organization."
            )
        return self.org.describe()

    def list_roots(self):
        return dict(
            Roots=[ou.describe() for ou in self.ou if isinstance(ou, FakeRoot)]
        )

    def create_organizational_unit(self, **kwargs):
        new_ou = FakeOrganizationalUnit(self.org, **kwargs)
        self.ou.append(new_ou)
        return new_ou.describe()

    def get_organizational_unit_by_id(self, ou_id):
        ou = next((ou for ou in self.ou if ou.id == ou_id), None)
        if ou is None:
            raise RESTError(
                'OrganizationalUnitNotFoundException',
                "You specified an organizational unit that doesn't exist."
            )
        return ou

    def validate_parent_id(self, parent_id):
        try:
            self.get_organizational_unit_by_id(parent_id)
        except RESTError:
            raise RESTError(
                'ParentNotFoundException',
                "You specified parent that doesn't exist."
            )
        return parent_id

    def describe_organizational_unit(self, **kwargs):
        ou = self.get_organizational_unit_by_id(kwargs['OrganizationalUnitId'])
        return ou.describe()

    def list_organizational_units_for_parent(self, **kwargs):
        parent_id = self.validate_parent_id(kwargs['ParentId'])
        return dict(
            OrganizationalUnits=[
                {
                    'Id': ou.id,
                    'Arn': ou.arn,
                    'Name': ou.name,
                }
                for ou in self.ou
                if ou.parent_id == parent_id
            ]
        )

    def create_account(self, **kwargs):
        new_account = FakeAccount(self.org, **kwargs)
        self.accounts.append(new_account)
        return new_account.create_account_status

    def get_account_by_id(self, account_id):
        account = next((
            account for account in self.accounts
            if account.id == account_id
        ), None)
        if account is None:
            raise RESTError(
                'AccountNotFoundException',
                "You specified an account that doesn't exist."
            )
        return account

    def describe_account(self, **kwargs):
        account = self.get_account_by_id(kwargs['AccountId'])
        return account.describe()

    def list_accounts(self):
        return dict(
            Accounts=[account.describe()['Account'] for account in self.accounts]
        )

    def list_accounts_for_parent(self, **kwargs):
        parent_id = self.validate_parent_id(kwargs['ParentId'])
        return dict(
            Accounts=[
                account.describe()['Account']
                for account in self.accounts
                if account.parent_id == parent_id
            ]
        )

    def move_account(self, **kwargs):
        new_parent_id = self.validate_parent_id(kwargs['DestinationParentId'])
        self.validate_parent_id(kwargs['SourceParentId'])
        account = self.get_account_by_id(kwargs['AccountId'])
        index = self.accounts.index(account)
        self.accounts[index].parent_id = new_parent_id

    def list_parents(self, **kwargs):
        if re.compile(r'[0-9]{12}').match(kwargs['ChildId']):
            child_object = self.get_account_by_id(kwargs['ChildId'])
        else:
            child_object = self.get_organizational_unit_by_id(kwargs['ChildId'])
        return dict(
            Parents=[
                {
                    'Id': ou.id,
                    'Type': ou.type,
                }
                for ou in self.ou
                if ou.id == child_object.parent_id
            ]
        )

    def list_children(self, **kwargs):
        parent_id = self.validate_parent_id(kwargs['ParentId'])
        if kwargs['ChildType'] == 'ACCOUNT':
            obj_list = self.accounts
        elif kwargs['ChildType'] == 'ORGANIZATIONAL_UNIT':
            obj_list = self.ou
        else:
            raise RESTError(
                'InvalidInputException',
                'You specified an invalid value.'
            )
        return dict(
            Children=[
                {
                    'Id': obj.id,
                    'Type': kwargs['ChildType'],
                }
                for obj in obj_list
                if obj.parent_id == parent_id
            ]
        )

    def create_policy(self, **kwargs):
        new_policy = FakeServiceControlPolicy(self.org, **kwargs)
        self.policies.append(new_policy)
        return new_policy.describe()

    def describe_policy(self, **kwargs):
        if re.compile(utils.SCP_ID_REGEX).match(kwargs['PolicyId']):
            policy = next((p for p in self.policies if p.id == kwargs['PolicyId']), None)
            if policy is None:
                raise RESTError(
                    'PolicyNotFoundException',
                    "You specified a policy that doesn't exist."
                )
        else:
            raise RESTError(
                'InvalidInputException',
                'You specified an invalid value.'
            )
        return policy.describe()

    def attach_policy(self, **kwargs):
        policy = next((p for p in self.policies if p.id == kwargs['PolicyId']), None)
        if (re.compile(utils.ROOT_ID_REGEX).match(kwargs['TargetId']) or
                re.compile(utils.OU_ID_REGEX).match(kwargs['TargetId'])):
            ou = next((ou for ou in self.ou if ou.id == kwargs['TargetId']), None)
            if ou is not None:
                if ou not in ou.attached_policies:
                    ou.attached_policies.append(policy)
                    policy.attachments.append(ou)
            else:
                raise RESTError(
                    'OrganizationalUnitNotFoundException',
                    "You specified an organizational unit that doesn't exist."
                )
        elif re.compile(utils.ACCOUNT_ID_REGEX).match(kwargs['TargetId']):
            account = next((a for a in self.accounts if a.id == kwargs['TargetId']), None)
            if account is not None:
                if account not in account.attached_policies:
                    account.attached_policies.append(policy)
                    policy.attachments.append(account)
            else:
                raise RESTError(
                    'AccountNotFoundException',
                    "You specified an account that doesn't exist."
                )
        else:
            raise RESTError(
                'InvalidInputException',
                'You specified an invalid value.'
            )

    def list_policies(self, **kwargs):
        return dict(Policies=[
            p.describe()['Policy']['PolicySummary'] for p in self.policies
        ])

    def list_policies_for_target(self, **kwargs):
        if re.compile(utils.OU_ID_REGEX).match(kwargs['TargetId']):
            obj = next((ou for ou in self.ou if ou.id == kwargs['TargetId']), None)
            if obj is None:
                raise RESTError(
                    'OrganizationalUnitNotFoundException',
                    "You specified an organizational unit that doesn't exist."
                )
        elif re.compile(utils.ACCOUNT_ID_REGEX).match(kwargs['TargetId']):
            obj = next((a for a in self.accounts if a.id == kwargs['TargetId']), None)
            if obj is None:
                raise RESTError(
                    'AccountNotFoundException',
                    "You specified an account that doesn't exist."
                )
        else:
            raise RESTError(
                'InvalidInputException',
                'You specified an invalid value.'
            )
        return dict(Policies=[
            p.describe()['Policy']['PolicySummary'] for p in obj.attached_policies
        ])

    def list_targets_for_policy(self, **kwargs):
        if re.compile(utils.SCP_ID_REGEX).match(kwargs['PolicyId']):
            policy = next((p for p in self.policies if p.id == kwargs['PolicyId']), None)
            if policy is None:
                raise RESTError(
                    'PolicyNotFoundException',
                    "You specified a policy that doesn't exist."
                )
        else:
            raise RESTError(
                'InvalidInputException',
                'You specified an invalid value.'
            )
        objects = [
            {
                'TargetId': obj.id,
                'Arn': obj.arn,
                'Name': obj.name,
                'Type': obj.type,
            } for obj in policy.attachments
        ]
        return dict(Targets=objects)


organizations_backend = OrganizationsBackend()