Merge pull request #173 from RiotGames/accounts-refactor

refactor: Refactored Account objects

backend/cloud_inquisitor/__init__.py

@@ -8,6 +8,7 @@
 from werkzeug.local import LocalProxy
 
 from cloud_inquisitor.constants import PLUGIN_NAMESPACES
+from cloud_inquisitor.exceptions import InquisitorError
 from cloud_inquisitor.utils import get_user_data_configuration, read_config
 
 logger = logging.getLogger(__name__)
@@ -45,6 +46,10 @@ def get_aws_session(account):
         :obj:`boto3:boto3.session.Session`
     """
     from cloud_inquisitor.config import dbconfig
+    from cloud_inquisitor.plugins.types.accounts import AWSAccount
+
+    if not isinstance(account, AWSAccount):
+        raise InquisitorError('Non AWSAccount  passed to get_aws_session, got {}'.format(account.__class__.__name__))
 
     # If no keys are on supplied for the account, use sts.assume_role instead
     session = get_local_aws_session()
@@ -103,6 +108,12 @@ def get_aws_regions(*, force=False):
     return __regions
 
 
+def get_plugin_by_name(ns, name):
+    for plugin in CINQ_PLUGINS[ns]['plugins']:
+        if plugin.name == name:
+            return plugin.load()
+
+
 # Check if the user has opted to use userdata based configuration for DB, and load it if needed
 if app_config.use_user_data:
     get_user_data_configuration()

backend/cloud_inquisitor/constants.py

@@ -18,6 +18,7 @@
     'notifier': 'cloud_inquisitor.plugins.notifiers',
     'schedulers': 'cloud_inquisitor.plugins.schedulers',
     'types': 'cloud_inquisitor.plugins.types',
+    'accounts': 'cloud_inquisitor.plugins.types.accounts',
     'view': 'cloud_inquisitor.plugins.views'
 })
 # endregion

backend/cloud_inquisitor/data/migrations/versions/cfb0ed4cced9_new_accounts_table.py

@@ -0,0 +1,195 @@
+"""New accounts table
+
+Revision ID: cfb0ed4cced9
+Revises: 1edcdfbc7780
+Create Date: 2018-07-10 13:26:01.588708
+
+"""
+from json import dumps as to_json
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy import text, inspect
+from sqlalchemy.dialects import mysql
+
+
+# revision identifiers, used by Alembic.
+revision = 'cfb0ed4cced9'
+down_revision = '1edcdfbc7780'
+
+select_ai = 'SELECT AUTO_INCREMENT FROM information_schema.TABLES WHERE TABLE_SCHEMA = :db AND TABLE_NAME = :table'
+select_cfg_item = 'SELECT value FROM config_items WHERE namespace_prefix = :ns AND `key` = :key'
+select_acct_types = 'SELECT account_type_id, account_type FROM account_types'
+insert_acct_type = 'INSERT INTO account_types (account_type) VALUES (:name)'
+insert_acct = (
+    'INSERT INTO accounts_new (account_id, account_name, account_type_id, contacts, enabled, required_roles)'
+    ' VALUES(:id, :name, :type_id, :contacts, :enabled, :required_roles)'
+)
+insert_acct_prop = 'INSERT INTO account_properties (account_id, name, value) VALUES (:id, :name, :value)'
+
+
+def upgrade():
+    create_new_tables()
+    migrate_data()
+    switch_tables()
+
+
+def downgrade():
+    raise Exception('You cannot downgrade from this version')
+
+
+def create_new_tables():
+    op.create_table('account_types',
+        sa.Column('account_type_id', mysql.INTEGER(unsigned=True), nullable=False, autoincrement=True),
+        sa.Column('account_type', sa.String(length=100), nullable=False),
+        sa.PrimaryKeyConstraint('account_type_id')
+    )
+    op.create_index(op.f('ix_account_types_account_type'), 'account_types', ['account_type'], unique=True)
+
+    op.create_table('accounts_new',
+        sa.Column('account_id', mysql.INTEGER(unsigned=True), nullable=False),
+        sa.Column('account_name', sa.String(length=256), nullable=False),
+        sa.Column('account_type_id', mysql.INTEGER(unsigned=True), nullable=False),
+        sa.Column('contacts', mysql.JSON(), nullable=False),
+        sa.Column('enabled', mysql.SMALLINT(unsigned=True), nullable=False),
+        sa.Column('required_roles', mysql.JSON(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ('account_type_id',),
+            ['account_types.account_type_id'],
+            name='fk_account_account_type_id',
+            ondelete='CASCADE'
+        ),
+        sa.PrimaryKeyConstraint('account_id')
+    )
+    op.create_index(op.f('ix_accounts_new_account_name'), 'accounts_new', ['account_name'], unique=True)
+    op.create_index(op.f('ix_accounts_new_account_type_id'), 'accounts_new', ['account_type_id'], unique=False)
+
+    op.create_table('account_properties',
+        sa.Column('property_id', mysql.INTEGER(unsigned=True), nullable=False, autoincrement=True),
+        sa.Column('account_id', mysql.INTEGER(unsigned=True), nullable=False),
+        sa.Column('name', sa.String(length=50), nullable=False),
+        sa.Column('value', mysql.JSON(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ('account_id',),
+            ['accounts_new.account_id'],
+            name='fk_account_properties_account_id',
+            ondelete='CASCADE'
+        ),
+        sa.PrimaryKeyConstraint('property_id', 'account_id')
+    )
+    op.create_index(op.f('ix_account_properties_account_id'), 'account_properties', ['account_id'], unique=False)
+    op.create_index(op.f('ix_account_properties_name'), 'account_properties', ['name'], unique=False)
+
+
+def migrate_data():
+    conn = op.get_bind()
+    account_types = {x['account_type']: x['account_type_id'] for x in conn.execute(text(select_acct_types))}
+    try:
+        schema = inspect(conn.engine).default_schema_name
+        conn.execute('SET FOREIGN_KEY_CHECKS=0')
+        res = conn.execute(text(select_ai), {'db': schema, 'table': 'accounts'})
+        acct_auto_increment = next(res)['AUTO_INCREMENT']
+
+        for acct_type in ('AWS', 'DNS: AXFR', 'DNS: CloudFlare'):
+            if acct_type not in account_types:
+                conn.execute(text(insert_acct_type), {'name': acct_type})
+                account_types[acct_type] = get_insert_id(conn)
+
+        res = conn.execute('SELECT * FROM accounts')
+        for acct in res:
+            if acct['account_type'] == 'AWS':
+                conn.execute(
+                    text(insert_acct),
+                    {
+                        'id': acct['account_id'],
+                        'name': acct['account_name'],
+                        'type_id': account_types['AWS'],
+                        'contacts': acct['contacts'],
+                        'enabled': acct['enabled'],
+                        'required_roles': acct['required_roles']
+                    }
+                )
+                conn.execute(
+                    text(insert_acct_prop),
+                    {
+                        'id': acct['account_id'],
+                        'name': 'account_number',
+                        'value': to_json(acct['account_number'])
+                    }
+                )
+
+                conn.execute(
+                    text(insert_acct_prop),
+                    {
+                        'id': acct['account_id'],
+                        'name': 'ad_group_base',
+                        'value': to_json(acct['ad_group_base'] or '')
+                    }
+                )
+                print('Migrated {} account {}'.format(acct['account_type'], acct['account_name']))
+
+            elif acct['account_type'] == 'DNS_AXFR':
+                conn.execute(
+                    text(insert_acct),
+                    {
+                        'id': acct['account_id'],
+                        'name': acct['account_name'],
+                        'type_id': account_types['DNS: AXFR'],
+                        'contacts': acct['contacts'],
+                        'enabled': acct['enabled'],
+                        'required_roles': acct['required_roles']
+                    }
+                )
+                server = get_config_value(conn, 'collector_dns', 'axfr_server')
+                domains = get_config_value(conn, 'collector_dns', 'axfr_domains')
+
+                conn.execute(text(insert_acct_prop), {'id': acct['account_id'], 'name': 'server', 'value': [server]})
+                conn.execute(text(insert_acct_prop), {'id': acct['account_id'], 'name': 'domains', 'value': domains})
+                print('Migrated {} account {}'.format(acct['account_type'], acct['account_name']))
+
+            elif acct['account_type'] == 'DNS_CLOUDFLARE':
+                conn.execute(
+                    text(insert_acct),
+                    {
+                        'id': acct['account_id'],
+                        'name': acct['account_name'],
+                        'type_id': account_types['DNS: CloudFlare'],
+                        'contacts': acct['contacts'],
+                        'enabled': acct['enabled'],
+                        'required_roles': acct['required_roles']
+                    }
+                )
+                api_key = get_config_value(conn, 'collector_dns', 'cloudflare_api_key')
+                email = get_config_value(conn, 'collector_dns', 'cloudflare_email')
+                endpoint = get_config_value(conn, 'collector_dns', 'cloudflare_endpoint')
+
+                conn.execute(text(insert_acct_prop), {'id': acct['account_id'], 'name': 'api_key', 'value': api_key})
+                conn.execute(text(insert_acct_prop), {'id': acct['account_id'], 'name': 'email', 'value': email})
+                conn.execute(text(insert_acct_prop), {'id': acct['account_id'], 'name': 'endpoint', 'value': endpoint})
+
+                print('Migrated {} account {}'.format(acct['account_type'], acct['account_name']))
+
+            else:
+                print('Invalid account type: {}'.format(acct['account_type']))
+        conn.execute(text('ALTER TABLE accounts_new AUTO_INCREMENT = :counter'), {'counter': acct_auto_increment})
+    finally:
+        conn.execute('SET FOREIGN_KEY_CHECKS=1')
+
+
+def switch_tables():
+    conn = op.get_bind()
+    conn.execute('SET FOREIGN_KEY_CHECKS=0')
+    conn.execute('DROP TABLE accounts')
+    conn.execute('ALTER TABLE resources MODIFY `account_id` int(10) unsigned')
+    conn.execute('ALTER TABLE accounts_new RENAME accounts')
+    conn.execute('ALTER TABLE accounts RENAME INDEX `ix_accounts_new_account_name` TO `ix_accounts_account_name`')
+    conn.execute('ALTER TABLE accounts RENAME INDEX `ix_accounts_new_account_type_id` TO `ix_accounts_account_type_id`')
+    conn.execute('SET FOREIGN_KEY_CHECKS=1')
+
+
+def get_insert_id(conn):
+    return next(conn.execute('SELECT LAST_INSERT_ID()'))[0]
+
+
+def get_config_value(conn, ns, item):
+    return next(conn.execute(text(select_cfg_item), {'ns': ns, 'key': item}))['value']

backend/cloud_inquisitor/exceptions.py

@@ -26,5 +26,9 @@ class IssueException(InquisitorError):
     """Exception class for issue types"""
 
 
+class AccountException(InquisitorError):
+    """Exception class for Account types"""
+
+
 class SchedulerError(InquisitorError):
     """Exception class for scheduler plugins"""

backend/cloud_inquisitor/plugins/commands/setup.py

@@ -1,51 +1,11 @@
-from click import confirm, prompt
-from flask_script import Option
-
 from cloud_inquisitor.app import initialize
-from cloud_inquisitor.database import db
 from cloud_inquisitor.plugins.commands import BaseCommand
-from cloud_inquisitor.schema import Account
 
 
 class Setup(BaseCommand):
     """Sets up the initial state of the configuration stored in the database"""
     name = 'Setup'
-    option_list = (
-        Option(
-            '-H', '--headless',
-            dest='headless_mode',
-            action='store_true',
-            help='Setup command will not prompt for interactive input'
-        ),
-    )
-
-    def init_account(self):
-        """Create a new account object"""
-        account_name = prompt('Account name', type=str)
-        account_number = prompt('Account ID Number', type=int)
-        contacts = prompt('Contacts', type=str)
-        if confirm('Limit access to this account?'):
-            required_roles = prompt('Required Role (optional)', type=str)
-        else:
-            required_roles = ''
-        enabled = confirm('Enabled')
-
-        acct = Account()
-        acct.account_name = account_name
-        acct.account_number = account_number
-        acct.contacts = list(map(lambda x: x.strip(), contacts.split(',')))
-        acct.enabled = enabled
-        acct.required_roles = [x for x in map(lambda x: x.strip(), required_roles.split(',')) if x]
-
-        db.session.add(acct)
-        db.session.commit()
-
-        self.log.info('Account has been added')
+    option_list = ()
 
     def run(self, **kwargs):
         initialize()
-
-        # If there are no accounts created, ask the user if he/she wants to create one now
-        if not kwargs['headless_mode'] and not db.Account.find_one():
-            if confirm('You have no accounts defined, do you wish to add the first account now?'):
-                self.init_account()

backend/cloud_inquisitor/plugins/commands/userdata.py

@@ -7,7 +7,7 @@
 
 from cloud_inquisitor import get_aws_session, app_config, get_local_aws_session
 from cloud_inquisitor.plugins.commands import BaseCommand
-from cloud_inquisitor.schema import Account
+from cloud_inquisitor.plugins.types.accounts import BaseAccount
 
 
 class UserData(BaseCommand):
@@ -55,7 +55,7 @@ def run(self, **kwargs):
                     print('you must set the kms_account_name setting in your configuration file to the name of the '
                           'account that is able to decrypt the user data')
                     return
-                acct = Account.get(kms_account_name)
+                acct = BaseAccount.get(kms_account_name)
                 if not acct:
                     print('You must add the {} account to the system for this to work'.format(kms_account_name))
                     return