Pluto/model.py

import os, sqlite3, json, urllib2, ssl, urllib, time, subprocess, socket

from flask import render_template_string
from util import *
import secrets

#db = sqlite3.connect(os.path.join(os.path.dirname(__file__), 'pluto.db'), check_same_thread = False)
db = sqlite3.connect('/var/www/pluto/pluto.db', check_same_thread = False)
cur = db.cursor()

so = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
so.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
so6 = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
so6.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)

class DBError(Exception):
    pass

class NoSuchEntity(DBError):
    pass

class TooManyEntities(DBError):
    pass

class DBObject(object):
    __FIELDS__ = ()
    __DEFAULTS__ = {}
    __TABLE__ = ''
    __TYPES__ = {}
    AUTO_COMMIT = True

    def __init__(self, rowid, *data):
        self.rowid = rowid
        for idx, field in enumerate(self.__FIELDS__):
            default = self.__DEFAULTS__.get(field)
            if idx < len(data):
                setattr(self, field, data[idx])
            else:
                setattr(self, field, default)

    @classmethod
    def create_table(cls):
        cur.execute('CREATE TABLE IF NOT EXISTS %(table)s (%(columns)s)'%\
                {'table': cls.__TABLE__,
                 'columns': ', '.join('%s%s'%(field, ' '+cls.__TYPES__[field] if field in cls.__TYPES__ else '') for field in cls.__FIELDS__)}
        )

    @classmethod
    def create(cls, *data):
        row = list(data)
        for field in cls.__FIELDS__[len(data):]:
            row.append(cls.__DEFAULTS__[field])
        cur.execute('INSERT INTO %(table)s VALUES (%(fields)s)'%{
            'table': cls.__TABLE__,
            'fields': ', '.join(['?'] * len(cls.__FIELDS__))
        }, row)
        if cls.AUTO_COMMIT:
            db.commit()
        return cls(cur.lastrowid, *row)

    def delete(self):
        cur.execute('DELETE FROM %(table)s WHERE ROWID=?'%{'table': self.__TABLE__}, (self.rowid,))
        if self.AUTO_COMMIT:
            db.commit()

    def update(self):
        cur.execute('UPDATE %(table)s SET %(fields)s WHERE ROWID=?'%{
            'table': self.__TABLE__,
            'fields': ', '.join('%s=?'%(field,) for field in self.__FIELDS__)
        }, tuple(getattr(self, field) for field in self.__FIELDS__) + (self.rowid,))
        if self.AUTO_COMMIT:
            db.commit()

    @classmethod
    def get(cls, **criteria):
        pairs = criteria.items()
        keys = [pair[0] for pair in pairs]
        values = [pair[1] for pair in pairs]
        cur.execute('SELECT ROWID, %(fields)s FROM %(table)s WHERE %(criteria)s'%{
            'table': cls.__TABLE__,
            'fields': ', '.join(cls.__FIELDS__),
            'criteria': ' and '.join('%s=?'%(k,) for k in keys),
        }, values)
        return [cls(*row) for row in cur]

    @classmethod
    def all(cls):
        cur.execute('SELECT ROWID, %(fields)s FROM %(table)s'%{
            'table': cls.__TABLE__,
            'fields': ', '.join(cls.__FIELDS__),
        })
        return [cls(*row) for row in cur]

    @classmethod
    def sorted(cls, by, limit=None):
        cur.execute('SELECT ROWID, %(fields)s FROM %(table)s ORDER BY %(by)s %(limit)s'%{
            'table': cls.__TABLE__,
            'fields': ', '.join(cls.__FIELDS__),
            'by': by,
            'limit': ('' if limit is None else 'LIMIT %d'%(limit,)),
        })
        return [cls(*row) for row in cur]

    @classmethod
    def get_one(cls, **criteria):
        res = cls.get(**criteria)
        if len(res) < 1:
            raise NoSuchEntity(cls, criteria)
        elif len(res) > 1:
            raise TooManyEntities(cls, criteria)
        return res[0]

    def __repr__(self):
        return '<%(cls)s(%(table)s %(row)d %(items)s'%{
                'table': self.__TABLE__,
                'cls': type(self).__name__,
                'row': self.rowid,
                'items': ' '.join('%s=%r'%(field, getattr(self, field)) for field in self.__FIELDS__),
        }


class Log(DBObject):
    __TABLE__ = 'log'
    __FIELDS__ = ('time', 'path', 'headers', 'data', 'hooks')

    @classmethod
    def most_recent(cls, n=None):
        return cls.sorted('time DESC', n)

class DebugLog(DBObject):
    __TABLE__ = 'debuglog'
    __FIELDS__ = ('time', 'path', 'headers', 'data', 'value', 'hook', 'cond', 'act', 'success', 'message')

    @classmethod
    def most_recent(cls, n=None):
        return cls.sorted('time DESC', n)

class Hook(DBObject):
    __TABLE__ = 'hooks'
    __FIELDS__ = ('name', 'author', 'disabled', 'debugged')
    __DEFAULTS__ = {
        'disabled': 0,
        'debugged': 0,
    }

    def trigger(self, path, headers, data, values, response):
        if self.disabled:
            return False
        conditions = Condition.for_hook(self)
        actions = Action.for_hook(self)
        for condition in conditions:
            result, msg = condition.test_select(path, headers, data, values, response)
            if self.debugged:
                DebugLog.create(time.time(), path, header_dumps(headers), jdumps(data), jdumps(values), self.rowid, condition.rowid, None, result, msg)
            if not result:
                break
        else:
            for act in actions:
                result = act.actuate(path, headers, data, values, response)
                if self.debugged:
                    DebugLog.create(time.time(), path, header_dumps(headers), jdumps(data), jdumps(values), self.rowid, None, act.rowid, None, result)
            if self.debugged:
                DebugLog.create(time.time(), path, header_dumps(headers), jdumps(data), jdumps(values), self.rowid, None, None, True, None)
            return True
        if self.debugged:
            DebugLog.create(time.time(), path, header_dumps(headers), jdumps(data), jdumps(values), self.rowid, None, None, False, None)
        return False
    
class Condition(DBObject):
    __TABLE__ = 'conditions'
    __FIELDS__ = ('hook', 'selector', 's1', 's2', 's3', 'test', 't1', 't2', 't3', 'invert')

    @classmethod
    def for_hook(cls, hook):
        return cls.get(hook=hook.rowid)

    def get_hook(self):
        return Hook.get_one(rowid=self.hook)

    def select(self, path, headers, data, values, response):
        return getattr(self, 'select_' + self.selector, self.no_select)(path, headers, data, values, response)

    def no_select(self, path, headers, data, values, response):
        print 'No selector found for', self.selector
        return None

    def select_header(self, path, headers, data, values, response):
        return headers.get(self.s1, '')

    def select_JSON(self, path, headers, data, values, response):
        if not isinstance(data, dict):
            return False
        cur = data
        for part in self.s1.split('.'):
            cur = cur.get(part)
            if cur is None:
                return False
        return str(cur)

    def select_path(self, path, headers, data, values, response):
        return path

    def select_value(self, path, headers, data, values, response):
        print values
        print self.s1
        print values.get(self.s1, '')
        return values.get(self.s1, '')

    def test_value(self, val):
        try:
            result = getattr(self, 'test_' + self.test, self.no_test)(val)
        except (ValueError, TypeError) as e:
            result = (False, "Error: " + str(e))
        if self.invert:
            result = (not result[0], result[1])
        return result

    def no_test(self, val):
        return False, "No valid test by that name"

    def test_equal(self, val):
        return str(val) == self.t1, "Compare: %r == %r" % (val, self.t1)

    def test_inrange(self, val):
        return float(self.t1) <= float(val) <= float(self.t2), "Compare %r <= %r <= %r" % (float(self.t1), float(val), float(self.t2))

    def test_truthy(self, val):
        return bool(val), "Test: %r" %(val,)

    def test_contains(self, val):
        return self.t1 in val, "Compare: %r in %r" % (self.t1, val)

    def test_select(self, path, headers, data, values, response):
        return self.test_value(self.select(path, headers, data, values, response))

class Action(DBObject):
    __TABLE__ = 'actions'
    __FIELDS__ = ('hook', 'action', 'a1', 'a2', 'a3')

    GITLAB_API = 'https://gitlab.cosi.clarkson.edu/api/v3/'
    GITLAB_TOKEN = secrets.GITLAB_TOKEN
    PROTO = ssl.PROTOCOL_TLSv1_2

    @classmethod
    def for_hook(cls, hook):
        return cls.get(hook=hook.rowid)

    def get_hook(self):
        return Hook.get_one(rowid=self.hook)

    def actuate(self, path, headers, data, values, response):
        try:
            return getattr(self, 'act_' + self.action, self.no_act)(path, headers, data, values, response)
        except (ValueError, TypeError):
            pass

    def no_act(self, path, headers, data, values, response):
        return 'INTERNAL ERROR: ACTION NOT FOUND'

    def act_post(self, path, headers, data, values, response):
        args = {'path': path, 'headers': headers, 'data': data, 'values': values}
        args['response'] = {'data': response.get_data(), 'headers': response.headers}
        url = render_template_string(self.a1, **args)
        postdata = render_template_string(self.a2, **args)
        headers = json.loads(render_template_string(self.a3, **args))
        print 'Note: posting to', url, 'with data', postdata, 'and headers', headers, '...'
        req = urllib2.Request(url, postdata, headers)
        ctxt = ssl.SSLContext(self.PROTO)
        result = urllib2.urlopen(req, context=ctxt)
        out = result.read()
        #out = None
        print 'Complete, got', repr(out)
        return out

    def act_gitlab(self, path, headers, data, values, response):
        args = {'path': path, 'headers': headers, 'data': data, 'values': values}
        args['response'] = {'data': response.get_data(), 'headers': response.headers}
        url = self.GITLAB_API + render_template_string(self.a1, **args)
        params = json.loads(render_template_string(self.a2, **args))
        headers = json.loads(render_template_string(self.a3, **args))
        headers.update({'PRIVATE-TOKEN': self.GITLAB_TOKEN})
        postdata = urllib.urlencode(params)
        print 'Note: posting to', url, 'with data', postdata, 'and headers', headers, '...'
        req = urllib2.Request(url, postdata, headers)
        ctxt = ssl.SSLContext(self.PROTO)
        result = urllib2.urlopen(req, context=ctxt)
        out = result.read()
        #out = None
        print 'Complete, got', repr(out)
        return out

    def act_system(self, path, headers, data, values, response):
        args = {'path': path, 'headers': headers, 'data': data, 'values': values}
        args['response'] = {'data': response.get_data(), 'headers': response.headers}
        cmd = render_template_string(self.a1, **args)
        if not self.a2:
			proc = subprocess.Popen(cmd, shell=True)
			return 'forked'
        else:
			try:
				return subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
			except subprocess.CalledProcessError as e:
				return e.output

    def act_udp(self, path, headers, data, values, response):
        args = {'path': path, 'headers': headers, 'data': data, 'values': values}
        args['response'] = {'data': response.get_data(), 'headers': response.headers}
        dest = render_template_string(self.a1, **args)
        packet = render_template_string(self.a2, **args)
        encoding = render_template_string(self.a3, **args)
        try:
            if encoding in (u'hex', u'base64'):
                packet = packet.decode(encoding)
            elif encoding == 'input':
                packet = str(data)
            elif encoding == 'json':
                packet = jdumps(data)  # XXX HACKS
            else:
                packet = packet.encode(encoding)
        except Exception as e:
            return 'failed to encode packet: ' + str(e)
        host, _, port = dest.partition(':')
        if not _:
            return 'illegal specification: no port in destination'
        try:
            port = int(port)
        except ValueError:
            return 'illegal port value: ' + port
        if port < 0 or port > 65535:
            return 'illegal port value: ' + str(port)
        try:
            res = socket.getaddrinfo(host, port)
        except socket.gaierror:
            return 'bad hostname:' + host
        for fam, tp, proto, canon, addr in res:
            if tp == socket.SOCK_DGRAM:
                try:
                    if fam == socket.AF_INET:
                        so.sendto(packet, addr)
                        return 'sent to {}: {}'.format(addr, packet.encode('hex'))
                    elif fam == socket.AF_INET6:
                        so6.sendto(packet, addr)
                        return 'sent to {}: {}'.format(addr, packet.encode('hex'))
                except Exception:
                    pass
        return 'no good address family found'

    def act_tcp(self, path, headers, data, values, response):
        args = {'path': path, 'headers': headers, 'data': data, 'values': values}
        args['response'] = {'data': response.get_data(), 'headers': response.headers}
        dest = render_template_string(self.a1, **args)
        packet = render_template_string(self.a2, **args)
        encoding = render_template_string(self.a3, **args)
        try:
            if encoding in (u'hex', u'base64'):
                packet = packet.decode(encoding)
            elif encoding == 'input':
                packet = str(data)
            elif encoding == 'json':
                packet = jdumps(data)  # XXX HACKS
            else:
                packet = packet.encode(encoding)
        except Exception as e:
            return 'failed to encode packet: ' + str(e)
        host, _, port = dest.partition(':')
        if not _:
            return 'illegal specification: no port in destination'
        try:
            port = int(port)
        except ValueError:
            return 'illegal port value: ' + port
        if port < 0 or port > 65535:
            return 'illegal port value: ' + str(port)
        try:
            res = socket.getaddrinfo(host, port)
        except socket.gaierror:
            return 'bad hostname:' + host
        so = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        so.settimeout(0.1)
        so6 = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        so6.settimeout(0.1)
        for fam, tp, proto, canon, addr in res:
            if tp == socket.SOCK_STREAM:
                try:
                    if fam == socket.AF_INET:
                        so.connect(addr)
                        so.send(packet)
                        return 'sent to {}: {}'.format(addr, packet.encode('hex'))
                    elif fam == socket.AF_INET6:
                        so6.connect(addr)
                        so6.send(packet)
                        return 'sent to {}: {}'.format(addr, packet.encode('hex'))
                except Exception:
                    pass
        return 'no good address family found'

    def act_set_response(self, path, headers, data, values, response):
        args = {'path': path, 'headers': headers, 'data': data, 'values': values}
        args['response'] = {'data': response.get_data(), 'headers': response.headers}
        content = render_template_string(self.a1, **args)
        content_type = render_template_string(self.a2, **args)
        response.set_data(content)
        response.headers['Content-type'] = content_type
        return 'response set to "' + content_type + '":\n' + content