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 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 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 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 TestConfig(Config):
    """
    This class specifies a configuration that is used for our tests.
    """
    TESTING = True
    DEBUG = True

    SQLALCHEMY_DATABASE_URI = 'postgresql://localhost:5432/doorman_test'

    WTF_CSRF_ENABLED = False

    DOORMAN_ENROLL_SECRET = [
        'secret',
    ]
    DOORMAN_EXPECTS_UNIQUE_HOST_ID = False

    DOORMAN_AUTH_METHOD = None

    DOORMAN_COLUMN_RENDER = {
        'computer_name': '<a href="https://{{ value | urlencode }}/">{{ value }}</a>'
    } def learn_from_result(result, node):
    if not result['data']:
        return

    capture_columns = set(
        map(itemgetter(0),
            current_app.config['DOORMAN_CAPTURE_NODE_INFO']
        )
    )

    if not capture_columns:
        return

    node_info = node.get('node_info', {})
    orig_node_info = node_info.copy()

    for _, action, columns, _, in extract_results(result):
        # only update columns common to both sets
        for column in capture_columns & set(columns):

            cvalue = node_info.get(column)  # current value
            value = columns.get(column)

            if action == 'removed' and (cvalue is None or cvalue != value):
                pass
            elif action == 'removed' and cvalue == value:
                node_info.pop(column)
            elif action == 'added' and (cvalue is None or cvalue != value):
                node_info[column] = value

    # only update node_info if there's actually a change

    if orig_node_info == node_info:
        return

    node = Node.get_by_id(node['id'])
    node.update(node_info=node_info)
    return class NodeFactory(BaseFactory):

    class Meta:
        model = Node class PackFactory(BaseFactory):

    class Meta:
        model = Pack class QueryFactory(BaseFactory):

    class Meta:
        model = Query class TagFactory(BaseFactory):

    class Meta:
        model = Tag
from copy import deepcopy
from flask import url_for
from sqlalchemy import and_
    from urlparse import urlparse
    from urllib.parse import urlparse
from doorman.models import (
    Node, Pack, Query, Tag, FilePath,
    DistributedQuery, DistributedQueryTask, DistributedQueryResult, Rule,
)
from doorman.settings import TestConfig
from doorman.utils import learn_from_result
from .factories import NodeFactory, PackFactory, QueryFactory, TagFactory
        from doorman.rules import Network
        from doorman.models import Rule
        from doorman.models import Rule
        from doorman.plugins import AbstractAlerterPlugin
        from doorman.rules import RuleMatch
        from doorman.tasks import analyze_result
import datetime as dt
import gzip
import io
import json
import mock
import time
        import doorman.utils
        import unicodecsv as csv
        q = DistributedQuery.create(
            sql="select name, path, pid from processes where name = 'osqueryd';")
        t = DistributedQueryTask.create(node=node, distributed_query=q)
        t.update(status=DistributedQueryTask.PENDING)

        data = [{
            "name": "osqueryd",
            "path": "/usr/local/bin/osqueryd",
            "pid": "97830"
        },
        {
            "name": "osqueryd",
            "path": "/usr/local/bin/osqueryd",
            "pid": "97831"
        }]

        r = DistributedQueryResult.create(columns=data[0],
                                          distributed_query=q,
                                          distributed_query_task=t)
        t.update(status=DistributedQueryTask.COMPLETE)

        resp = testapp.post_json(url_for('api.distributed_write'), {
            'node_key': node.node_key,
            'queries': {
                t.guid: '',
            }
        },
        extra_environ=dict(REMOTE_ADDR='127.0.0.2')
        )

        assert q.results
        assert len(q.results) == 1
        assert q.results[0] == r
        assert q.results[0].columns == data[0]
        assert node.last_ip == '127.0.0.2'

    def test_distributed_query_write_state_failed(self, db, node, testapp):
        q = DistributedQuery.create(
            sql="select name, path, pid from processes where name = 'osqueryd';")
        t = DistributedQueryTask.create(node=node, distributed_query=q)
        t.update(status=DistributedQueryTask.PENDING)

        data = []

        resp = testapp.post_json(url_for('api.distributed_write'), {
            'node_key': node.node_key,
            'queries': {
                t.guid: data,
            },
            'statuses': {
                t.guid: 2,
            }
        },
        extra_environ=dict(REMOTE_ADDR='127.0.0.2')
        )

        assert t.status == DistributedQueryTask.FAILED
        assert not q.results
        assert node.last_ip == '127.0.0.2'

    def test_malicious_node_distributed_query_write(self, db, node, testapp):
        foo = NodeFactory(host_identifier='foo')
        q1 = DistributedQuery.create(
            sql="select name, path, pid from processes where name = 'osqueryd';")
        t1 = DistributedQueryTask.create(node=node, distributed_query=q1)
        q2 = DistributedQuery.create(
            sql="select name, path, pid from processes where name = 'osqueryd';")
        t2 = DistributedQueryTask.create(node=foo, distributed_query=q2)

        t1.update(status=DistributedQueryTask.PENDING)
        t2.update(status=DistributedQueryTask.PENDING)

        resp = testapp.post_json(url_for('api.distributed_write'), {
            'node_key': foo.node_key,
            'queries': {
                t1.guid: 'bar'
            }
        },
        extra_environ=dict(REMOTE_ADDR='127.0.0.2')
        )

        assert not q1.results
        assert not q2.results
        assert node.last_ip != '127.0.0.2'
        assert not node.last_ip
        assert foo.last_ip == '127.0.0.2'

        resp = testapp.post_json(url_for('api.distributed_write'), {
            'node_key': foo.node_key,
            'queries': {
                t2.guid: 'bar'
            }
        },
        extra_environ=dict(REMOTE_ADDR='127.0.0.3')
        )

        assert t2.results
        assert node.last_ip != '127.0.0.2'
        assert not node.last_ip
        assert foo.last_ip == '127.0.0.3'


class TestDistributedTable:

    def html_escape(self, v):
        return (v
            .replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
            .replace("'", "&#39;").replace('"', "&quot;")
            )

    def test_distributed_query_table_basic(self, db, node, testapp):
        # Create two fake queries, tasks, and fake results.
        q1 = DistributedQuery.create(
            sql="select name, path, pid from processes where name = 'osqueryd';")
        t1 = DistributedQueryTask.create(node=node, distributed_query=q1)
        q2 = DistributedQuery.create(
            sql="select name, path, pid from processes where name = 'otherproc';")
        t2 = DistributedQueryTask.create(node=node, distributed_query=q2)

        t1.update(status=DistributedQueryTask.PENDING)

        t2.update(status=DistributedQueryTask.PENDING)