def shell_cmd(verbose, with_req_context):
try:
from IPython.terminal.ipapp import TerminalIPythonApp
except ImportError:
click.echo(
cformat(
'%{red!}You need to `pip install ipython` to use the SNMS shell'
))
sys.exit(1)
current_app.config['REPL'] = True # disables e.g. memoize_request
request_stats_request_started()
context, info = _make_shell_context()
banner = cformat('%{yellow!}SNMS v{} is ready for your commands').format(
snms.__version__)
if verbose:
banner = '\n'.join(info + ['', banner])
ctx = current_app.make_shell_context()
ctx.update(context)
# clearCache()
stack = ExitStack()
if with_req_context:
stack.enter_context(
current_app.test_request_context(base_url=config.BASE_URL))
with stack:
ipython_app = TerminalIPythonApp.instance(user_ns=ctx,
display_banner=False)
ipython_app.initialize(argv=[])
ipython_app.shell.show_banner(banner)
ipython_app.start()
def create(grant_admin, create_company):
"""Creates a new user"""
user_type = 'user' if not grant_admin else 'super_admin'
while True:
email = prompt_email()
if email is None:
return
email = email.lower()
if not User.query.filter(User.email == email).count():
break
print(cformat('%{red}Email already exists'))
name = click.prompt("Name").strip()
print()
password = prompt_pass()
if password is None:
return
user = User(uid=generate_uid(), email=email, name=name, role=user_type, password=generate_password_hash(password), is_verified = True)
if create_company:
company_name = click.prompt("Company Name").strip()
company = Company(name=company_name, uid=generate_uid(), key=generate_key())
user_company = CompanyUserAssociation(role=ROLE_ADMIN)
user_company.user = user
company.users.append(user_company)
user.own_companies.append(company)
db.session.add(company)
_print_user_info(user)
if click.confirm(cformat("%{yellow}Create the new {}?").format(user_type), default=True):
db.session.add(user)
db.session.commit()
print(cformat("%{green}New {} created successfully with ID: %{green!}{}").format(user_type, user.uid))
def celery_cmd(args):
# remove the celery shell command
next(funcs for group, funcs, _ in command_classes if group == 'Main').remove('shell')
del CeleryCommand.commands['shell']
if args and args[0] == 'flower':
# Somehow flower hangs when executing it using CeleryCommand() so we simply exec it directly.
# It doesn't really need the celery config anyway (besides the broker url)
try:
import flower
except ImportError:
print(cformat('%{red!}Flower is not installed'))
sys.exit(1)
auth_args = ['[email protected]', '--auth_provider=flower.views.auth.GithubLoginHandler']
auth_env = {'FLOWER_OAUTH2_KEY': '',
'FLOWER_OAUTH2_SECRET': '',
'FLOWER_OAUTH2_REDIRECT_URI': 'http://localhost:5555/login'
}
args = ['celery', '-b', config.CELERY_BROKER] + args
env = dict(os.environ, **auth_env)
os.execvpe('celery', args, env)
elif args and args[0] == 'shell':
print(cformat('%{red!}Please use `snms shell`.'))
sys.exit(1)
else:
CeleryCommand(celery).execute_from_commandline(['snms celery'] + args)
def _require_encoding(encoding):
cur_encoding = db.engine.execute("SELECT current_setting('server_encoding')").scalar()
if cur_encoding >= encoding:
return True
print(cformat('%{red}Database encoding must be {}; got {}').format(encoding, cur_encoding))
print(cformat('%{yellow}Recreate your database using `createdb -E {} -T template0 ...`').format(encoding))
return False
def _require_extensions(*names):
missing = sorted(name for name in names if not has_extension(db.engine, name))
if not missing:
return True
print(cformat('%{red}Required Postgres extensions missing: {}').format(', '.join(missing)))
print(cformat('%{yellow}Create them using these SQL commands (as a Postgres superuser):'))
for name in missing:
print(cformat('%{white!} CREATE EXTENSION {};').format(name))
return False
def create_all_tables(db, verbose=False, add_initial_data=True):
"""Create all tables and required initial objects"""
from snms.modules.settings.operations import add_default_settings
if verbose:
print(cformat('%{green}Creating tables'))
db.create_all()
if verbose:
print(cformat('%{green}Tables, created'))
if add_initial_data:
if verbose:
print(cformat('%{green}Adding default settings'))
add_default_settings()
def purge():
"""Remove deleted companies, sensors, and other data."""
print(cformat('%{yellow!}*** DANGER'))
print(cformat('%{yellow!}***%{reset} '
'%{red!}This operation will %{yellow!}PERMANENTLY ERASE %{red!} data!%{reset}'))
from snms.modules.sensors import Sensor, SensorType
deleted_sensors = Sensor.find(Sensor.deleted).all()
tabledata = [['ID', 'Name', 'Created On']]
for sensor in deleted_sensors:
tabledata.append([str(sensor.id), sensor.name, sensor.created_at])
table = AsciiTable(tabledata, cformat('%{white!}Deleted Sensors%{reset}'))
print(table.table)
def _add_to_context(namespace,
info,
element,
name=None,
doc=None,
color='green'):
if not name:
name = element.__name__
namespace[name] = element
if doc:
info.append(
cformat('+ %%{%s}{}%%{white!} ({})' % color).format(name, doc))
else:
info.append(cformat('+ %%{%s}{}' % color).format(name))
def revoke_admin(user_id):
"""Revokes administration rights from a given user"""
user = User.get(user_id)
if user is None:
print(cformat("%{red}This user does not exist"))
return
_print_user_info(user)
if not user.is_admin:
print(cformat("%{yellow}This user does not have administration rights"))
return
if click.confirm(cformat("%{yellow}Revoke administration rights from this user?")):
user.role = 'user'
db.session.commit()
print(cformat("%{green}Administration rights revoked successfully"))
def grant_admin(user_id):
"""Grants administration rights to a given user"""
user = User.get(user_id)
if user is None:
print(cformat("%{red}This user does not exist"))
return
_print_user_info(user)
if user.role == ROLE_ADMIN:
print(cformat("%{yellow}This user already has administration rights"))
return
if click.confirm(cformat("%{yellow}Grant administration rights to this user?")):
user.role = ROLE_ADMIN
db.session.commit()
print(cformat("%{green}Administration rights granted successfully"))
def _add_to_context_multi(namespace,
info,
elements,
names=None,
doc=None,
color='green'):
if not names:
names = [x.__name__ for x in elements]
for name, element in zip(names, elements):
namespace[name] = element
if doc:
info.append(
cformat('+ %%{white!}{}:%%{reset} %%{%s}{}' % color).format(
doc, ', '.join(names)))
else:
info.append(cformat('+ %%{%s}{}' % color).format(', '.join(names)))
def _require_pg_version(version):
# convert version string such as '9.4.10' to `90410` which is the
# format used by server_version_num
req_version = sum(segment * 10**(4 - 2*i) for i, segment in enumerate(map(int, version.split('.'))))
cur_version = db.engine.execute("SELECT current_setting('server_version_num')::int").scalar()
if cur_version >= req_version:
return True
print(cformat('%{red}Postgres version too old; you need at least {} (or newer)').format(version))
return False
def prepare_db(empty=False, root_path=None, verbose=True):
"""Initialize an empty database (create tables, set alembic rev to HEAD)."""
if not _require_pg_version('9.2'):
return
# if not _require_encoding('UTF8'):
# return
# if not _require_extensions('unaccent', 'pg_trgm'):
# return
root_path = root_path or current_app.root_path
tables = get_all_tables(db)
if 'alembic_version' not in tables['public']:
if verbose:
print(cformat('%{green}Setting the alembic version to HEAD'))
stamp(directory=os.path.join(root_path, 'migrations'), revision='heads')
PluginScriptDirectory.dir = os.path.join(root_path, 'core', 'plugins', 'alembic')
alembic.command.ScriptDirectory = PluginScriptDirectory
plugin_msg = cformat("%{cyan}Setting the alembic version of the %{cyan!}{}%{reset}%{cyan} "
"plugin to HEAD%{reset}")
# for plugin in plugin_engine.get_active_plugins().values():
# if not os.path.exists(plugin.alembic_versions_path):
# continue
# if verbose:
# print(plugin_msg.format(plugin.name))
# with plugin.plugin_context():
# stamp(revision='heads')
# Retrieve the table list again, just in case we created unexpected tables
tables = get_all_tables(db)
tables['public'] = [t for t in tables['public'] if not t.startswith('alembic_version')]
if any(tables.values()):
if verbose:
print(cformat('%{red}Your database is not empty!'))
print(cformat('%{yellow}If you just added a new table/model, create an alembic revision instead!'))
print()
print('Tables in your database:')
for schema, schema_tables in sorted(tables.items()):
for t in schema_tables:
print(cformat(' * %{cyan}{}%{reset}.%{cyan!}{}%{reset}').format(schema, t))
return
create_all_tables(db, verbose=verbose, add_initial_data=(not empty))
def _make_shell_context():
context = {}
info = [cformat('%{white!}Available objects')]
add_to_context = partial(_add_to_context, context, info)
add_to_context_multi = partial(_add_to_context_multi, context, info)
add_to_context_smart = partial(_add_to_context_smart, context, info)
# Common stdlib modules
info.append(cformat('*** %{magenta!}stdlib%{reset} ***'))
DATETIME_ATTRS = ('date', 'time', 'datetime', 'timedelta')
ORM_ATTRS = ('joinedload', 'defaultload', 'contains_eager', 'lazyload',
'noload', 'subqueryload', 'undefer', 'undefer_group',
'load_only')
add_to_context_multi([getattr(datetime, attr) for attr in DATETIME_ATTRS] +
[getattr(sqlalchemy.orm, attr)
for attr in ORM_ATTRS] + [itertools, re, sys, os],
color='yellow')
# Models
info.append(cformat('*** %{magenta!}Models%{reset} ***'))
models = [
cls for name, cls in sorted(db.Model._decl_class_registry.items(),
key=itemgetter(0))
if hasattr(cls, '__table__')
]
add_to_context_smart(models)
# Tasks
info.append(cformat('*** %{magenta!}Tasks%{reset} ***'))
tasks = [
task for task in celery.tasks.values()
if not task.name.startswith('celery.')
]
add_to_context_smart(tasks,
get_name=lambda x: x.name.replace('.', '_'),
color='blue!')
# Plugins
info.append(cformat('*** %{magenta!}Plugins%{reset} ***'))
# Utils
info.append(cformat('*** %{magenta!}Misc%{reset} ***'))
add_to_context(celery, 'celery', doc='celery app', color='blue!')
add_to_context(db, 'db', doc='sqlalchemy db interface', color='cyan!')
add_to_context(now_utc,
'now_utc',
doc='get current utc time',
color='cyan!')
add_to_context(config, 'config', doc='snms config')
add_to_context(current_app, 'app', doc='flask app')
add_to_context(lambda *a, **kw: server_to_utc(datetime.datetime(*a, **kw)),
'dt',
doc='like datetime() but converted from localtime to utc')
# Stuff from plugins
# signals.plugin.shell_context.send(add_to_context=add_to_context, add_to_context_multi=add_to_context_multi)
return context, info
def _safe_downgrade(*args, **kwargs):
func = kwargs.pop('_func')
print(cformat('%{yellow!}*** DANGER'))
print(cformat('%{yellow!}***%{reset} '
'%{red!}This operation may %{yellow!}PERMANENTLY ERASE %{red!}some data!%{reset}'))
if current_app.debug:
skip_confirm = os.environ.get('SNMS_ALWAYS_DOWNGRADE', '').lower() in ('1', 'yes')
print(cformat('%{yellow!}***%{reset} '
"%{green!}Debug mode is active, so you probably won't destroy valuable data"))
else:
skip_confirm = False
print(cformat('%{yellow!}***%{reset} '
"%{red!}Debug mode is NOT ACTIVE, so make sure you are on the right machine!"))
if not skip_confirm and input(cformat('%{yellow!}***%{reset} '
'To confirm this, enter %{yellow!}YES%{reset}: ')) != 'YES':
print(cformat('%{green}Aborted%{reset}'))
sys.exit(1)
else:
return func(*args, **kwargs)
