class AddDistributedQueryForm(Form):

    sql = TextAreaField('Query', validators=[DataRequired(), ValidSQL()])
    description = TextAreaField('Description', validators=[Optional()])
    not_before = DateTimeField('Not Before', format="%Y-%m-%d %H:%M:%S",
                               validators=[Optional()])
    nodes = SelectMultipleField('Nodes', choices=[])
    tags = SelectMultipleField('Tags', choices=[])

    def set_choices(self):
        from doorman.models import Node, Tag
        self.nodes.choices = Node.query.with_entities(Node.node_key, Node.host_identifier).all()
        self.tags.choices = Tag.query.with_entities(Tag.value, Tag.value).all() class CreateQueryForm(QueryForm):

    def validate(self):
        from doorman.models import Query
        initial_validation = super(CreateQueryForm, self).validate()
        if not initial_validation:
            return False

        query = Query.query.filter(Query.name == self.name.data).first()
        if query:
            self.name.errors.append(
                u"Query with the name {0} already exists!".format(
                self.name.data)
            )
            return False

        # TODO could do some validation of the sql query
        return True class UpdateQueryForm(QueryForm):

    def __init__(self, *args, **kwargs):
        super(UpdateQueryForm, self).__init__(*args, **kwargs)
        self.set_choices()
        query = kwargs.pop('obj', None)
        if query:
            self.packs.process_data([p.name for p in query.packs])
            self.tags.process_data('\n'.join(t.value for t in query.tags)) class CreateTagForm(Form):
    value = TextAreaField('Tag', validators=[DataRequired()]) class UploadPackForm(Form):

    pack = FileField(u'Pack configuration', validators=[FileRequired()]) class FilePathForm(Form):
    category = StringField('category', validators=[DataRequired()])
    target_paths = TextAreaField('files', validators=[DataRequired()])
    tags = TextAreaField("Tags") class FilePathUpdateForm(FilePathForm):

    def __init__(self, *args, **kwargs):
        super(FilePathUpdateForm, self).__init__(*args, **kwargs)
        # self.set_choices()
        file_path = kwargs.pop('obj', None)
        if file_path:
            self.target_paths.process_data('\n'.join(file_path.get_paths()))
            self.tags.process_data('\n'.join(t.value for t in file_path.tags)) class CreateRuleForm(RuleForm):

    def validate(self):
        from doorman.models import Rule

        initial_validation = super(CreateRuleForm, self).validate()
        if not initial_validation:
            return False

        query = Rule.query.filter(Rule.name == self.name.data).first()
        if query:
            self.name.errors.append(
                u"Rule with the name {0} already exists!".format(
                self.name.data)
            )
            return False

        return True class UpdateRuleForm(RuleForm):

    def __init__(self, *args, **kwargs):
        super(UpdateRuleForm, self).__init__(*args, **kwargs)
        self.set_choices() class UpdateNodeForm(Form):

    display_name = StringField('Name', validators=[Optional()])
    is_active = BooleanField('Active', validators=[Optional()]) class CRUDMixin(object):
class Model(CRUDMixin, db.Model):
class SurrogatePK(object):
    def create(cls, **kwargs):
    def update(self, commit=True, **kwargs):
    def save(self, commit=True):
    def delete(self, commit=True):
    def get_by_id(cls, record_id):
def reference_col(tablename, nullable=False, pk_name='id', **kwargs): class DistributedQuery(SurrogatePK, Model):

    description = Column(db.String, nullable=True)
    sql = Column(db.String, nullable=False)
    timestamp = Column(db.DateTime, default=dt.datetime.utcnow)
    not_before = Column(db.DateTime, default=dt.datetime.utcnow)

    def __init___(self, sql, description=None, not_before=None):
        self.sql = sql
        self.description = description
        self.not_before = not_before class DistributedQueryTask(SurrogatePK, Model):

    NEW = 0
    PENDING = 1
    COMPLETE = 2
    FAILED = 3

    guid = Column(db.String, nullable=False, unique=True)
    status = Column(db.Integer, default=0, nullable=False)
    timestamp = Column(db.DateTime)

    distributed_query_id = reference_col('distributed_query', nullable=False)
    distributed_query = relationship(
        'DistributedQuery',
        backref=db.backref('tasks',
                           cascade='all, delete-orphan',
                           lazy='dynamic'),
    )

    node_id = reference_col('node', nullable=False)
    node = relationship(
        'Node',
        backref=db.backref('distributed_queries', lazy='dynamic'),
    )

    def __init__(self, node=None, node_id=None,
                 distributed_query=None, distributed_query_id=None):
        self.guid = str(uuid.uuid4())
        if node:
            self.node = node
        elif node_id:
            self.node_id = node_id
        if distributed_query:
            self.distributed_query = distributed_query
        elif distributed_query_id:
            self.distributed_query_id = distributed_query_id

    @declared_attr
    def __table_args__(cls):
        return (
            Index('idx_%s_node_id_status' % cls.__tablename__, 'node_id', 'status'),
        ) class DistributedQueryResult(SurrogatePK, Model):

    columns = Column(JSONB)
    timestamp = Column(db.DateTime, default=dt.datetime.utcnow)

    distributed_query_task_id = reference_col('distributed_query_task', nullable=False)
    distributed_query_task = relationship(
        'DistributedQueryTask',
        backref=db.backref('results',
                           cascade='all, delete-orphan',
                           lazy='joined'),
    )

    distributed_query_id = reference_col('distributed_query', nullable=False)
    distributed_query = relationship(
        'DistributedQuery',
        backref=db.backref('results',
                           cascade='all, delete-orphan',
                           lazy='joined'),
    )

    def __init__(self, columns, distributed_query=None, distributed_query_task=None):
        self.columns = columns
        self.distributed_query = distributed_query
        self.distributed_query_task = distributed_query_task class FilePath(SurrogatePK, Model):

    category = Column(db.String, nullable=False, unique=True)
    target_paths = Column(db.String)

    tags = relationship(
        'Tag',
        secondary=file_path_tags,
        back_populates='file_paths',
        lazy='joined',
    )

    def __init__(self, category=None, target_paths=None, *args, **kwargs):
        self.category = category

        if target_paths is not None:
            self.set_paths(*target_paths)
        elif args:
            self.set_paths(*args)
        else:
            self.target_paths = ''

    def to_dict(self):
        return {
            self.category: self.get_paths()
        }

    def get_paths(self):
        return self.target_paths.split('!!')

    def set_paths(self, *target_paths):
        self.target_paths = '!!'.join(target_paths) class Node(SurrogatePK, Model):

    node_key = Column(db.String, nullable=False, unique=True)
    enroll_secret = Column(db.String)
    enrolled_on = Column(db.DateTime)
    host_identifier = Column(db.String)
    last_checkin = Column(db.DateTime)
    node_info = Column(JSONB, default={}, nullable=False)
    is_active = Column(db.Boolean, default=True, nullable=False)
    last_ip = Column(INET, nullable=True)

    tags = relationship(
        'Tag',
        secondary=node_tags,
        back_populates='nodes',
        lazy='joined',
    )

    def __init__(self, host_identifier, node_key=None,
                 enroll_secret=None, enrolled_on=None, last_checkin=None,
                 is_active=True, last_ip=None,
                 **kwargs):
        self.node_key = node_key or str(uuid.uuid4())
        self.host_identifier = host_identifier
        self.enroll_secret = enroll_secret
        self.enrolled_on = enrolled_on
        self.last_checkin = last_checkin
        self.is_active = is_active
        self.last_ip = last_ip

    def __repr__(self):
        return '<Node-{0.id}: node_key={0.node_key}, host_identifier={0.host_identifier}>'.format(self)

    def get_config(self, **kwargs):
        from doorman.utils import assemble_configuration
        return assemble_configuration(self)

    def get_new_queries(self, **kwargs):
        from doorman.utils import assemble_distributed_queries
        return assemble_distributed_queries(self)

    @property
    def display_name(self):
        if 'display_name' in self.node_info and self.node_info['display_name']:
            return self.node_info['display_name']
        elif 'hostname' in self.node_info and self.node_info['hostname']:
            return self.node_info['hostname']
        elif 'computer_name' in self.node_info and self.node_info['computer_name']:
            return self.node_info['computer_name']
        else:
            return self.host_identifier

    @property
    def packs(self):
        return db.session.object_session(self) \
            .query(Pack) \
            .join(pack_tags, pack_tags.c['pack.id'] == Pack.id) \
            .join(node_tags, node_tags.c['tag.id'] == pack_tags.c['tag.id']) \
            .filter(node_tags.c['node.id'] == self.id) \
            .options(db.lazyload('*'))

    @property
    def queries(self):
        return db.session.object_session(self) \
            .query(Query) \
            .join(query_tags, query_tags.c['query.id'] == Query.id) \
            .join(node_tags, node_tags.c['tag.id'] == query_tags.c['tag.id']) \
            .filter(node_tags.c['node.id'] == self.id) \
            .options(db.lazyload('*'))

    @property
    def file_paths(self):
        return db.session.object_session(self) \
            .query(FilePath) \
            .join(file_path_tags, file_path_tags.c['file_path.id'] == FilePath.id) \
            .join(node_tags, node_tags.c['tag.id'] == file_path_tags.c['tag.id']) \
            .filter(node_tags.c['node.id'] == self.id) \
            .options(db.lazyload('*'))

    def to_dict(self):
        # NOTE: deliberately not including any secret values in here, for now.
        return {
            'id': self.id,
            'display_name': self.display_name,
            'enrolled_on': self.enrolled_on,
            'host_identifier': self.host_identifier,
            'last_checkin': self.last_checkin,
            'node_info': self.node_info.copy(),
            'last_ip': self.last_ip,
            'is_active': self.is_active
        } class Pack(SurrogatePK, Model):

    name = Column(db.String, nullable=False, unique=True)
    platform = Column(db.String)
    version = Column(db.String)
    description = Column(db.String)
    shard = Column(db.Integer)

    queries = relationship(
        'Query',
        secondary=querypacks,
        back_populates='packs',
    )

    tags = relationship(
        'Tag',
        secondary=pack_tags,
        back_populates='packs',
    )

    def __init__(self, name, platform=None, version=None,
                 description=None, shard=None, **kwargs):
        self.name = name
        self.platform = platform
        self.version = version
        self.description = description
        self.shard = shard

    def __repr__(self):
        return '<Pack: {0.name}>'.format(self)

    def to_dict(self):
        queries = {}
        discovery = []

        for query in self.queries:
            if 'discovery' in (t.value for t in query.tags):
                discovery.append(query.sql)
            else:
                queries[query.name] = query.to_dict()

        return {
            'platform': self.platform,
            'version': self.version,
            'shard': self.shard,
            'discovery': discovery,
            'queries': queries,
        } class Query(SurrogatePK, Model):

    name = Column(db.String, nullable=False)
    sql = Column(db.String, nullable=False)
    interval = Column(db.Integer, default=3600)
    platform = Column(db.String)
    version = Column(db.String)
    description = Column(db.String)
    value = Column(db.String)
    removed = Column(db.Boolean, nullable=False, default=True)
    shard = Column(db.Integer)

    packs = relationship(
        'Pack',
        secondary=querypacks,
        back_populates='queries',
    )

    tags = relationship(
        'Tag',
        secondary=query_tags,
        back_populates='queries',
        lazy='joined',
    )

    def __init__(self, name, query=None, sql=None, interval=3600, platform=None,
                 version=None, description=None, value=None, removed=True,
                 shard=None, **kwargs):
        self.name = name
        self.sql = query or sql
        self.interval = int(interval)
        self.platform = platform
        self.version = version
        self.description = description
        self.value = value
        self.removed = removed
        self.shard = shard

    def __repr__(self):
        return '<Query: {0.name}>'.format(self)

    def to_dict(self):
        return {
            'query': self.sql,
            'interval': self.interval,
            'platform': self.platform,
            'version': self.version,
            'description': self.description,
            'value': self.value,
            'removed': self.removed,
            'shard': self.shard,
        } class Tag(SurrogatePK, Model):

    value = Column(db.String, nullable=False, unique=True)

    nodes = relationship(
        'Node',
        secondary=node_tags,
        back_populates='tags',
    )

    packs = relationship(
        'Pack',
        secondary=pack_tags,
        back_populates='tags',
    )

    queries = relationship(
        'Query',
        secondary=query_tags,
        back_populates='tags',
    )

    file_paths = relationship(
        'FilePath',
        secondary=file_path_tags,
        back_populates='tags',
    )

    def __init__(self, value, **kwargs):
        self.value = value

    def __repr__(self):
        return '<Tag: {0.value}>'.format(self)

    @property
    def packs_count(self):
        return db.session.object_session(self) \
            .query(Pack.id).with_parent(self, 'packs').count()

    @property
    def nodes_count(self):
        return db.session.object_session(self) \
            .query(Node.id).with_parent(self, 'nodes').count()

    @property
    def queries_count(self):
        return db.session.object_session(self) \
            .query(Query.id).with_parent(self, 'queries').count()

    @property
    def file_paths_count(self):
        return db.session.object_session(self) \
            .query(FilePath.id).with_parent(self, 'file_paths').count() class Rule(SurrogatePK, Model):

    name = Column(db.String, nullable=False)
    alerters = Column(ARRAY(db.String), nullable=False)
    description = Column(db.String, nullable=True)
    conditions = Column(JSONB)
    updated_at = Column(db.DateTime, nullable=False, default=dt.datetime.utcnow)

    def __init__(self, name, alerters, description=None, conditions=None, updated_at=None):
        self.name = name
        self.description = description
        self.alerters = alerters
        self.conditions = conditions
        self.updated_at = updated_at

    @property
    def template(self):
        return string.Template("{name}\r\n\r\n{description}".format(
            name=self.name, description=self.description or '')
        ) class ResultLog(SurrogatePK, Model):

    name = Column(db.String, nullable=False)
    timestamp = Column(db.DateTime, default=dt.datetime.utcnow)
    action = Column(db.String)
    columns = Column(JSONB)

    node_id = reference_col('node', nullable=False)
    node = relationship(
        'Node',
        backref=db.backref('result_logs', lazy='dynamic')
    )

    def __init__(self, name=None, action=None, columns=None, timestamp=None,
                 node=None, node_id=None, **kwargs):
        self.name = name
        self.action = action
        self.columns = columns or {}
        self.timestamp = timestamp
        if node:
            self.node = node
        elif node_id:
            self.node_id = node_id

    @declared_attr
    def __table_args__(cls):
        return (
            Index('idx_%s_node_id_timestamp_desc' % cls.__tablename__,
                  'node_id', cls.timestamp.desc()),
        ) class StatusLog(SurrogatePK, Model):

    line = Column(db.Integer)
    message = Column(db.String)
    severity = Column(db.Integer)
    filename = Column(db.String)
    created = Column(db.DateTime, default=dt.datetime.utcnow)
    version = Column(db.String)

    node_id = reference_col('node', nullable=False)
    node = relationship(
        'Node',
        backref=db.backref('status_logs', lazy='dynamic')
    )

    def __init__(self, line=None, message=None, severity=None,
                 filename=None, created=None, node=None, node_id=None,
                 version=None, **kwargs):
        self.line = int(line)
        self.message = message
        self.severity = int(severity)
        self.filename = filename
        self.created = created
        self.version = version
        if node:
            self.node = node
        elif node_id:
            self.node_id = node_id

    @declared_attr
    def __table_args__(cls):
        return (
            Index('idx_%s_node_id_created_desc' % cls.__tablename__,
                'node_id', cls.created.desc()),
        ) def create_query_pack_from_upload(upload):
    '''
    Create a pack and queries from a query pack file. **Note**, if a
    pack already exists under the filename being uploaded, then any
    queries defined here will be added to the existing pack! However,
    if a query with a particular name already exists, and its sql is
    NOT the same, then a new query with the same name but different id
    will be created (as to avoid clobbering the existing query). If its
    sql is identical, then the query will be reused.

    '''
    # The json package on Python 3 expects a `str` input, so we're going to
    # read the body and possibly convert to the right type
    body = upload.data.read()
    if not isinstance(body, six.string_types):
        body = body.decode('utf-8')

    try:
        data = json.loads(body)
    except ValueError:
        flash(u"Could not load pack as JSON - ensure it is JSON encoded",
              'danger')
        return None
    else:
        if 'queries' not in data:
            flash(u"No queries in pack", 'danger')
            return None

        name = splitext(basename(upload.data.filename))[0]
        pack = Pack.query.filter(Pack.name == name).first()

    if not pack:
        current_app.logger.debug("Creating pack %s", name)
        pack = Pack.create(name=name, **data)

    for query_name, query in data['queries'].items():
        if not validate_osquery_query(query['query']):
            flash('Invalid osquery query: "{0}"'.format(query['query']), 'danger')
            return None

        q = Query.query.filter(Query.name == query_name).first()

        if not q:
            q = Query.create(name=query_name, **query)
            pack.queries.append(q)
            current_app.logger.debug("Adding new query %s to pack %s",
                                     q.name, pack.name)
            continue

        if q in pack.queries:
            continue

        if q.sql == query['query']:
            current_app.logger.debug("Adding existing query %s to pack %s",
                                     q.name, pack.name)
            pack.queries.append(q)
        else:
            q2 = Query.create(name=query_name, **query)
            current_app.logger.debug(
                "Created another query named %s, but different sql: %r vs %r",
                query_name, q2.sql.encode('utf-8'), q.sql.encode('utf-8'))
            pack.queries.append(q2)

    else:
        pack.save()
        flash(u"Imported query pack {0}".format(pack.name), 'success')

    return pack def flash_errors(form):
    '''http://flask.pocoo.org/snippets/12/'''
    for field, errors in form.errors.items():
        for error in errors:
            message = u"Error in the {0} field - {1}".format(
                getattr(form, field).label.text, error
            )
            flash(message, 'danger') def get_paginate_options(request, model, choices, existing_query=None,
                         default='id', page=1, max_pp=20, default_sort='asc'):

    try:
        per_page = int(request.args.get('pp', max_pp))
    except Exception:
        per_page = 20

    per_page = max(0, min(max_pp, per_page))

    order_by = request.args.get('order_by', 'id')
    if order_by not in choices:
        order_by = default
    order_by = getattr(model, order_by, 'id')

    sort = request.args.get('sort', default_sort)
    if sort not in ('asc', 'desc'):
        sort = default_sort

    order_by = getattr(order_by, sort)()

    if existing_query:
        query = existing_query.order_by(order_by)
    else:
        query = model.query.order_by(order_by)

    query = query.paginate(page=page, per_page=per_page)
    return query
from io import BytesIO
from operator import itemgetter
from flask import (
    Blueprint, current_app, flash, jsonify, redirect, render_template,
    request, send_file, url_for
)
from flask_login import login_required
from flask_paginate import Pagination
from sqlalchemy import or_
from sqlalchemy.orm import joinedload
from .forms import (
    AddDistributedQueryForm,
    CreateQueryForm,
    UpdateQueryForm,
    CreateTagForm,
    UploadPackForm,
    FilePathForm,
    FilePathUpdateForm,
    CreateRuleForm,
    UpdateRuleForm,
    UpdateNodeForm,
)
from doorman.database import db
from doorman.models import (
    DistributedQuery, DistributedQueryTask, DistributedQueryResult,
    FilePath, Node, Pack, Query, Tag, Rule, ResultLog, StatusLog
)
from doorman.utils import (
    create_query_pack_from_upload, flash_errors, get_paginate_options
)
import json
import datetime as dt
import unicodecsv as csv
@blueprint.route('/nodes.csv')
@login_required
def nodes_csv():
    headers = [
        'Display name',
        'Host identifier',
        'Enrolled On',
        'Last Check-in',
        'Last IP Address',
        'Is Active',
    ]

    column_names = map(itemgetter(0), current_app.config['DOORMAN_CAPTURE_NODE_INFO'])
    labels = map(itemgetter(1), current_app.config['DOORMAN_CAPTURE_NODE_INFO'])
    headers.extend(labels)
    headers = list(map(str.title, headers))

    bio = BytesIO()
    writer = csv.writer(bio)
    writer.writerow(headers)

    for node in Node.query:
        row = [
            node.display_name,
            node.host_identifier,
            node.enrolled_on,
            node.last_checkin,
            node.last_ip,
            node.is_active,
        ]
        row.extend([node.node_info.get(column, '') for column in column_names])
        writer.writerow(row)

    bio.seek(0)

    response = send_file(
        bio,
        mimetype='text/csv',
        as_attachment=True,
        attachment_filename='nodes.csv'
    )

    return response


@blueprint.route('/nodes/add', methods=['GET', 'POST'])
@login_required
def add_node():
    return redirect(url_for('manage.nodes'))


@blueprint.route('/nodes/tagged/<string:tags>')
@login_required
def nodes_by_tag(tags):
    if tags == 'null':
        nodes = Node.query.filter(Node.tags == None).all()
    else:
        tag_names = [t.strip() for t in tags.split(',')]
        nodes = Node.query.filter(Node.tags.any(Tag.value.in_(tag_names))).all()
    return render_template('nodes.html', nodes=nodes)


@blueprint.route('/node/<int:node_id>', methods=['GET', 'POST'])
@login_required
def get_node(node_id):
    node = Node.query.filter_by(id=node_id).first_or_404()
    form = UpdateNodeForm(request.form)

    if form.validate_on_submit():
        node_info = node.node_info.copy()

        if form.display_name.data:
            node_info['display_name'] = form.display_name.data
        elif 'display_name' in node_info:
            node_info.pop('display_name')

        node.node_info = node_info
        node.is_active = form.is_active.data
        node.save()

        if request.is_xhr:
            return '', 204

        return redirect(url_for('manage.get_node', node_id=node.id))

    form = UpdateNodeForm(request.form, obj=node)
    flash_errors(form)

    packs = node.packs \
        .options(
            db.joinedload(Pack.tags, innerjoin=True),
            db.joinedload(Pack.queries, innerjoin=True),
        ).order_by(Pack.name)

    queries = node.queries \
        .options(
            db.joinedload(Query.tags, innerjoin=True),
            db.joinedload(Query.packs)
        ).order_by(Query.name)

    return render_template('node.html', form=form, node=node,
                           packs=packs, queries=queries)


@blueprint.route('/node/<int:node_id>/activity')
@login_required
def node_activity(node_id):
    node = Node.query.filter_by(id=node_id) \
        .options(db.lazyload('*')).first()

    try:
        timestamp = request.args.get('timestamp')
        timestamp = dt.datetime.fromtimestamp(float(timestamp))
    except Exception:
        timestamp = dt.datetime.utcnow()
        timestamp -= dt.timedelta(days=7)

    recent = node.result_logs.filter(ResultLog.timestamp > timestamp).all()
    queries = db.session.query(DistributedQueryTask) \
        .join(DistributedQuery) \

        .join(DistributedQueryResult) \