<commit_before>import functools

def vmap(f, dic):
    return dict(zip(dic.keys(), map(f, dic.values())))

class Formula(object):
    @property
    def __name__(self):
        return self.get_lambda().__name__

    def __init__(self):
        self.is_cache = True
        self.cache_lambda = None
        self.cache_consts = None
        self.children = []

    def __str__(self):
        return self.create_lambda_string()

    def debug(self):
        return (self.create_lambda_string(), list(self.traverse_const_values()))

    def nocache(self):
        self.is_cache = False
        return self

    def create_lambda_string(self):
        traversed = list(self.traverse())
        body = "".join(traversed)

        constized_args = sorted(list(set(self.traverse_constize_args())))
        unnamed_args = list(self.traverse_args())
        if len(constized_args) == 0:
            args = ",".join(unnamed_args)
        elif len(unnamed_args) == 0:
            args = ",".join(constized_args)
        else:
            raise SyntaxError("_ and _1 ~ _9 can not be used at the same time.")
        return "lambda {}:{}".format(args, body)

    def create_lambda(self):
        binds = self.traverse_const_values()
        lambda_string = self.create_lambda_string()
        return eval(lambda_string, dict(binds))

    def get_lambda(self):
        if not self.is_cache or self.cache_lambda is None:
            self.cache_lambda = self.create_lambda()
        return self.cache_lambda

    def __call__(self, *args):
        return self.get_lambda()(*args)

    @classmethod
    def convert_oprand(cls, x):
        if not issubclass(x.__class__, Formula):
            return ConstOperand(x)
        if isinstance(x, Underscore) and x.id == 0:
            return Underscore()
        return x

    def do_operator2(self, other, operator):
        this = Formula.convert_oprand(self)
        other = Formula.convert_oprand(other)
        return Operator2(operator, this, other)

    def rdo_operator2(self, other, operator):
        this = Formula.convert_oprand(self)
        other = Formula.convert_oprand(other)
        return Operator2(operator, other, this)

    def do_operator1(self, operator):
        this = Formula.convert_oprand(self)
        return Operator1(operator, this)

    def do_getitem(self, item):
        this = Formula.convert_oprand(self)
        item = Formula.convert_oprand(item)
        return GetItem(this, item)

    def do_methodcall(self, method):
        def f(*args, **kwargs):
            this = Formula.convert_oprand(self)
            return MethodCall(
                this,
                method,
                list(map(Formula.convert_oprand, args)),
                vmap(Formula.convert_oprand, kwargs)
            )
        return f

    def __add__(self, other):
        return self.do_operator2(other, "+")

    def __sub__(self, other):
        return self.do_operator2(other, "-")

    def __mul__(self, other):
        return self.do_operator2(other, "*")

    def __div__(self, other):
        return self.do_operator2(other, "/")

    def __truediv__(self, other):
        return self.do_operator2(other, "/")

    def __floordiv__(self, other):
        return self.do_operator2(other, "//")

    def __mod__(self, other):
        return self.do_operator2(other, "%")

    def __lt__(self, other):
        return self.do_operator2(other, "<")

    def __le__(self, other):
        return self.do_operator2(other, "<=")

    def __gt__(self, other):
        return self.do_operator2(other, ">")

    def __ge__(self, other):
        return self.do_operator2(other, ">=")

    def __eq__(self, other):
        return self.do_operator2(other, "==")

    def __ne__(self, other):
        return self.do_operator2(other, "!=")

    def __pow__(self, other):
        return self.do_operator2(other, "**")

    def __lshift__(self, other):
        return self.do_operator2(other, "<<")

    def __rshift__(self, other):
        return self.do_operator2(other, ">>")

    def __and__(self, other):
        return self.do_operator2(other, "&")

    def __or__(self, other):
        return self.do_operator2(other, "|")

    def __xor__(self, other):
        return self.do_operator2(other, "^")

    def __radd__(self, other):
        return self.rdo_operator2(other, "+")

    def __rsub__(self, other):
        return self.rdo_operator2(other, "-")

    def __rmul__(self, other):
        return self.rdo_operator2(other, "*")

    def __rdiv__(self, other):
        return self.rdo_operator2(other, "/")

    def __rtruediv__(self, other):
        return self.rdo_operator2(other, "/")

    def __rfloordiv__(self, other):
        return self.rdo_operator2(other, "//")

    def __rmod__(self, other):
        return self.rdo_operator2(other, "%")

    def __rlt__(self, other):
        return self.rdo_operator2(other, "<")

    def __rle__(self, other):
        return self.rdo_operator2(other, "<=")

    def __rgt__(self, other):
        return self.rdo_operator2(other, ">")

    def __rge__(self, other):
        return self.rdo_operator2(other, ">=")

    def __req__(self, other):
        return self.rdo_operator2(other, "==")

    def __rne__(self, other):
        return self.rdo_operator2(other, "!=")

    def __rpow__(self, other):
        return self.rdo_operator2(other, "**")

    def __rlshift__(self, other):
        return self.rdo_operator2(other, "<<")

    def __rrshift__(self, other):
        return self.rdo_operator2(other, ">>")

    def __rand__(self, other):
        return self.rdo_operator2(other, "&")

    def __ror__(self, other):
        return self.rdo_operator2(other, "|")

    def __rxor__(self, other):
        return self.rdo_operator2(other, "^")

    def __pos__(self):
        return self.do_operator1("+")

    def __neg__(self):
        return self.do_operator1("-")

    def __invert__(self):
        return self.do_operator1("~")

    def __getattr__(self, method):
        return self.do_methodcall(method)

    def __getitem__(self, key):
        return self.do_getitem(key)

    def traverse_const_values(self):
        for child in self.children:
            for t in child.traverse_const_values():
                yield t

    def traverse_args(self):
        for child in self.children:
            for t in child.traverse_args():
                yield t

    def traverse_constize_args(self):
        for child in self.children:
            for t in child.traverse_constize_args():
                yield t

class Operator(Formula):
    pass

class Operator1(Operator):
    def __init__(self, operator, value):
        super(Operator1, self).__init__()
        self.operator = operator
        self.value = value
        self.children = [self.value]

    def traverse(self):
        yield '('
        yield self.operator
        for t in self.value.traverse():
            yield t
        yield ')'

class Operator2(Operator):
    def __init__(self, operator, left, right):
        super(Operator2, self).__init__()
        self.operator = operator
        self.left = left
        self.right = right
        self.children = [self.left, self.right]

    def traverse(self):
        yield '('
        for t in self.left.traverse():
            yield t
        yield self.operator
        for t in self.right.traverse():
            yield t
        yield ')'

class Operand(Formula):
    pass

class ConstOperand(Operand):
    COUNTER = 0

    def __init__(self, value):
        super(ConstOperand, self).__init__()
        self.id = ConstOperand.COUNTER
        ConstOperand.COUNTER += 1
        self.value = value

    def traverse(self):
        yield '('
        yield "CONST_{}".format(self.id)
        yield ')'

    def traverse_const_values(self):
        yield ('CONST_{}'.format(self.id), self.value)

class Underscore(Operand):
    NUMBER_CONSTIZE = 10
    COUNTER = NUMBER_CONSTIZE

    def __init__(self, id=None):
        super(Underscore, self).__init__()
        if id is None:
            self.id = Underscore.COUNTER
            Underscore.COUNTER += 1
        else:
            self.id = id

    def traverse(self):
        yield '('
        yield "___ARG{}___".format(self.id)
        yield ')'

    def traverse_args(self):
        if self.id >= Underscore.NUMBER_CONSTIZE:
            yield "___ARG{}___".format(self.id)

    def traverse_constize_args(self):
        if self.id < Underscore.NUMBER_CONSTIZE:
            yield "___ARG{}___".format(self.id)

class MethodCall(Formula):
    def __init__(self, value, method, args, kwargs):
        super(MethodCall, self).__init__()
        self.value = value
        self.method = method
        self.args = args
        self.kwargs = kwargs
        self.children = [self.value] + self.args + list(self.kwargs.values())

    def traverse(self):
        yield '('
        for t in self.value.traverse():
            yield t
        yield '.'
        yield self.method
        yield '('
        for arg in self.args:
            for t in arg.traverse():
                yield t
            yield ','
        for name, arg in self.kwargs:
            yield '{}='.format(name)
            for t in self.arg.traverse():
                yield t
            yield ','
        yield ')'
        yield ')'

class GetItem(Formula):
    def __init__(self, value, item):
        super(GetItem, self).__init__()
        self.value = value
        self.item = item
        self.children = [self.value, self.item]

    def traverse(self):
        yield '('
        for t in self.value.traverse():
            yield t
        yield '['
        for t in self.item.traverse():
            yield t
        yield ']'
        yield ')'

class FunctionCall(Formula):
    COUNTER = 0

    def __init__(self, func, args, kwargs):
        super(FunctionCall, self).__init__()
        self.id = FunctionCall.COUNTER
        FunctionCall.COUNTER += 1
        self.func = func
        self.args = args
        self.kwargs = kwargs
        self.children = self.args + list(self.kwargs.values())

    def traverse(self):
        yield "BIND_FUNC_{}".format(self.id)
        yield '('
        for arg in self.args:
            for t in arg.traverse():
                yield t
            yield ','
        for name, arg in self.kwargs:
            yield '{}='.format(name)
            for t in arg.traverse():
                yield t
            yield ','
        yield ')'

    def traverse_const_values(self):
        yield ('BIND_FUNC_{}'.format(self.id), self.func)
        for t in super(FunctionCall, self).traverse_const_values():
            yield t

def scalambdable_func(*funcs):
    @functools.wraps(funcs[0])
    def wraps(*args, **kwargs):
        for f in reversed(funcs):
            def is_scalambda_object(x):
                return issubclass(x.__class__, Formula)
            if any(map(is_scalambda_object, args)) or any(map(is_scalambda_object, kwargs.values())):
                args = [FunctionCall(f, list(map(Formula.convert_oprand, args)), vmap(Formula.convert_oprand, kwargs))]
                kwargs = {}
            else:
                args = [f(*args, **kwargs)]
                kwargs = {}
        return args[0]
    return wraps

_ = Underscore(0)
_1 = Underscore(1)
_2 = Underscore(2)
_3 = Underscore(3)
_4 = Underscore(4)
_5 = Underscore(5)
_6 = Underscore(6)
_7 = Underscore(7)
_8 = Underscore(8)
_9 = Underscore(9)
SF = scalambdable_func
<commit_msg>もし、strにしてevalしても値が変わらないならそのまま文字列としてlambda式の中に埋め込むことで、変数を参照しに行くコストを下げる。<commit_after>import functools

def vmap(f, dic):
    return dict(zip(dic.keys(), map(f, dic.values())))

def can_str_emmbed(value):
    try:
        return eval(str(value)) == value
    except:
        return False

class Formula(object):
    @property
    def __name__(self):
        return self.get_lambda().__name__

    def __init__(self):
        self.is_cache = True
        self.cache_lambda = None
        self.cache_consts = None
        self.children = []

    def __str__(self):
        return self.create_lambda_string()

    def debug(self):
        return (self.create_lambda_string(), list(self.traverse_const_values()))

    def nocache(self):
        self.is_cache = False
        return self

    def create_lambda_string(self):
        traversed = list(self.traverse())
        body = "".join(traversed)

        constized_args = sorted(list(set(self.traverse_constize_args())))
        unnamed_args = list(self.traverse_args())
        if len(constized_args) == 0:
            args = ",".join(unnamed_args)
        elif len(unnamed_args) == 0:
            args = ",".join(constized_args)
        else:
            raise SyntaxError("_ and _1 ~ _9 can not be used at the same time.")
        return "lambda {}:{}".format(args, body)

    def create_lambda(self):
        binds = self.traverse_const_values()
        lambda_string = self.create_lambda_string()
        return eval(lambda_string, dict(binds))

    def get_lambda(self):
        if not self.is_cache or self.cache_lambda is None:
            self.cache_lambda = self.create_lambda()
        return self.cache_lambda

    def __call__(self, *args):
        return self.get_lambda()(*args)

    @classmethod
    def convert_oprand(cls, x):
        if not issubclass(x.__class__, Formula):
            return ConstOperand(x)
        if isinstance(x, Underscore) and x.id == 0:
            return Underscore()
        return x

    def do_operator2(self, other, operator):
        this = Formula.convert_oprand(self)
        other = Formula.convert_oprand(other)
        return Operator2(operator, this, other)

    def rdo_operator2(self, other, operator):
        this = Formula.convert_oprand(self)
        other = Formula.convert_oprand(other)
        return Operator2(operator, other, this)

    def do_operator1(self, operator):
        this = Formula.convert_oprand(self)
        return Operator1(operator, this)

    def do_getitem(self, item):
        this = Formula.convert_oprand(self)
        item = Formula.convert_oprand(item)
        return GetItem(this, item)

    def do_methodcall(self, method):
        def f(*args, **kwargs):
            this = Formula.convert_oprand(self)
            return MethodCall(
                this,
                method,
                list(map(Formula.convert_oprand, args)),
                vmap(Formula.convert_oprand, kwargs)
            )
        return f

    def __add__(self, other):
        return self.do_operator2(other, "+")

    def __sub__(self, other):
        return self.do_operator2(other, "-")

    def __mul__(self, other):
        return self.do_operator2(other, "*")

    def __div__(self, other):
        return self.do_operator2(other, "/")

    def __truediv__(self, other):
        return self.do_operator2(other, "/")

    def __floordiv__(self, other):
        return self.do_operator2(other, "//")

    def __mod__(self, other):
        return self.do_operator2(other, "%")

    def __lt__(self, other):
        return self.do_operator2(other, "<")

    def __le__(self, other):
        return self.do_operator2(other, "<=")

    def __gt__(self, other):
        return self.do_operator2(other, ">")

    def __ge__(self, other):
        return self.do_operator2(other, ">=")

    def __eq__(self, other):
        return self.do_operator2(other, "==")

    def __ne__(self, other):
        return self.do_operator2(other, "!=")

    def __pow__(self, other):
        return self.do_operator2(other, "**")

    def __lshift__(self, other):
        return self.do_operator2(other, "<<")

    def __rshift__(self, other):
        return self.do_operator2(other, ">>")

    def __and__(self, other):
        return self.do_operator2(other, "&")

    def __or__(self, other):
        return self.do_operator2(other, "|")

    def __xor__(self, other):
        return self.do_operator2(other, "^")

    def __radd__(self, other):
        return self.rdo_operator2(other, "+")

    def __rsub__(self, other):
        return self.rdo_operator2(other, "-")

    def __rmul__(self, other):
        return self.rdo_operator2(other, "*")

    def __rdiv__(self, other):
        return self.rdo_operator2(other, "/")

    def __rtruediv__(self, other):
        return self.rdo_operator2(other, "/")

    def __rfloordiv__(self, other):
        return self.rdo_operator2(other, "//")

    def __rmod__(self, other):
        return self.rdo_operator2(other, "%")

    def __rlt__(self, other):
        return self.rdo_operator2(other, "<")

    def __rle__(self, other):
        return self.rdo_operator2(other, "<=")

    def __rgt__(self, other):
        return self.rdo_operator2(other, ">")

    def __rge__(self, other):
        return self.rdo_operator2(other, ">=")

    def __req__(self, other):
        return self.rdo_operator2(other, "==")

    def __rne__(self, other):
        return self.rdo_operator2(other, "!=")

    def __rpow__(self, other):
        return self.rdo_operator2(other, "**")

    def __rlshift__(self, other):
        return self.rdo_operator2(other, "<<")

    def __rrshift__(self, other):
        return self.rdo_operator2(other, ">>")

    def __rand__(self, other):
        return self.rdo_operator2(other, "&")

    def __ror__(self, other):
        return self.rdo_operator2(other, "|")

    def __rxor__(self, other):
        return self.rdo_operator2(other, "^")

    def __pos__(self):
        return self.do_operator1("+")

    def __neg__(self):
        return self.do_operator1("-")

    def __invert__(self):
        return self.do_operator1("~")

    def __getattr__(self, method):
        return self.do_methodcall(method)

    def __getitem__(self, key):
        return self.do_getitem(key)

    def traverse_const_values(self):
        for child in self.children:
            for t in child.traverse_const_values():
                yield t

    def traverse_args(self):
        for child in self.children:
            for t in child.traverse_args():
                yield t

    def traverse_constize_args(self):
        for child in self.children:
            for t in child.traverse_constize_args():
                yield t

class Operator(Formula):
    pass

class Operator1(Operator):
    def __init__(self, operator, value):
        super(Operator1, self).__init__()
        self.operator = operator
        self.value = value
        self.children = [self.value]

    def traverse(self):
        yield '('
        yield self.operator
        for t in self.value.traverse():
            yield t
        yield ')'

class Operator2(Operator):
    def __init__(self, operator, left, right):
        super(Operator2, self).__init__()
        self.operator = operator
        self.left = left
        self.right = right
        self.children = [self.left, self.right]

    def traverse(self):
        yield '('
        for t in self.left.traverse():
            yield t
        yield self.operator
        for t in self.right.traverse():
            yield t
        yield ')'

class Operand(Formula):
    pass

class ConstOperand(Operand):
    COUNTER = 0

    def __init__(self, value):
        super(ConstOperand, self).__init__()
        self.id = ConstOperand.COUNTER
        ConstOperand.COUNTER += 1
        self.value = value

        self.is_use_dict = not can_str_emmbed(self.value)

    def traverse(self):
        if self.is_use_dict:
            yield '('
            yield "CONST_{}".format(self.id)
            yield ')'
        else:
            yield '('
            yield str(self.value)
            yield ')'

    def traverse_const_values(self):
        if self.is_use_dict:
            yield ('CONST_{}'.format(self.id), self.value)

class Underscore(Operand):
    NUMBER_CONSTIZE = 10
    COUNTER = NUMBER_CONSTIZE

    def __init__(self, id=None):
        super(Underscore, self).__init__()
        if id is None:
            self.id = Underscore.COUNTER
            Underscore.COUNTER += 1
        else:
            self.id = id

    def traverse(self):
        yield '('
        yield "___ARG{}___".format(self.id)
        yield ')'

    def traverse_args(self):
        if self.id >= Underscore.NUMBER_CONSTIZE:
            yield "___ARG{}___".format(self.id)

    def traverse_constize_args(self):
        if self.id < Underscore.NUMBER_CONSTIZE:
            yield "___ARG{}___".format(self.id)

class MethodCall(Formula):
    def __init__(self, value, method, args, kwargs):
        super(MethodCall, self).__init__()
        self.value = value
        self.method = method
        self.args = args
        self.kwargs = kwargs
        self.children = [self.value] + self.args + list(self.kwargs.values())

    def traverse(self):
        yield '('
        for t in self.value.traverse():
            yield t
        yield '.'
        yield self.method
        yield '('
        for arg in self.args:
            for t in arg.traverse():
                yield t
            yield ','
        for name, arg in self.kwargs:
            yield '{}='.format(name)
            for t in self.arg.traverse():
                yield t
            yield ','
        yield ')'
        yield ')'

class GetItem(Formula):
    def __init__(self, value, item):
        super(GetItem, self).__init__()
        self.value = value
        self.item = item
        self.children = [self.value, self.item]

    def traverse(self):
        yield '('
        for t in self.value.traverse():
            yield t
        yield '['
        for t in self.item.traverse():
            yield t
        yield ']'
        yield ')'

class FunctionCall(Formula):
    COUNTER = 0

    def __init__(self, func, args, kwargs):
        super(FunctionCall, self).__init__()
        self.id = FunctionCall.COUNTER
        FunctionCall.COUNTER += 1
        self.func = func
        self.args = args
        self.kwargs = kwargs
        self.children = self.args + list(self.kwargs.values())

    def traverse(self):
        yield "BIND_FUNC_{}".format(self.id)
        yield '('
        for arg in self.args:
            for t in arg.traverse():
                yield t
            yield ','
        for name, arg in self.kwargs:
            yield '{}='.format(name)
            for t in arg.traverse():
                yield t
            yield ','
        yield ')'

    def traverse_const_values(self):
        yield ('BIND_FUNC_{}'.format(self.id), self.func)
        for t in super(FunctionCall, self).traverse_const_values():
            yield t

def scalambdable_func(*funcs):
    @functools.wraps(funcs[0])
    def wraps(*args, **kwargs):
        for f in reversed(funcs):
            def is_scalambda_object(x):
                return issubclass(x.__class__, Formula)
            if any(map(is_scalambda_object, args)) or any(map(is_scalambda_object, kwargs.values())):
                args = [FunctionCall(f, list(map(Formula.convert_oprand, args)), vmap(Formula.convert_oprand, kwargs))]
                kwargs = {}
            else:
                args = [f(*args, **kwargs)]
                kwargs = {}
        return args[0]
    return wraps

_ = Underscore(0)
_1 = Underscore(1)
_2 = Underscore(2)
_3 = Underscore(3)
_4 = Underscore(4)
_5 = Underscore(5)
_6 = Underscore(6)
_7 = Underscore(7)
_8 = Underscore(8)
_9 = Underscore(9)
SF = scalambdable_func
<|endoftext|>


<commit_before>#!/usr/bin/env python3

###############################################################################
# Module Imports
###############################################################################

import arrow
import itertools
import logging
import shutil

from bs4 import BeautifulSoup as bs4
from collections import namedtuple
from lxml import etree, html
from pkgutil import get_data
from pyscp import SnapshotConnector, use_default_logging
from tempfile import TemporaryDirectory
from pathlib import Path

###############################################################################
# Global Constants And Variables
###############################################################################

log = logging.getLogger('pyscp.ebook')

###############################################################################
# Primary Classes
###############################################################################

class Epub:

    item = namedtuple('EbookItem', 'uid children')

    def __init__(self, title):
        self.dir = TemporaryDirectory()
        self.title = title
        self.lang = 'en'
        self.root = []
        self.spine = []
        self.uid = ('{:04}'.format(i) for i in itertools.count(1))

    def _template(self, name):
        name_to_file = dict(page='page.xhtml', spine='content.opf',
                            container='container.xml', toc='toc.ncx')
        return etree.ElementTree(etree.fromstring(
            get_data('pyscp', 'resources/templates/' + name_to_file[name])))

    def _write_to_disk(self, uid, title, content):
        page_xml = self._template('page')
        page_xml.xpath("/*/*[1]/*[1]")[0].text = title
        page_xml.xpath("/*/*[2]")[0].append(html.fromstring(content))
        page_xml.write("{}/{}.xhtml".format(self.dir.name, uid))

    def add_page(self, title, content, parent=None):
        item = self.item(next(self.uid), [])
        if not parent:
            self.root.append(item)
        else:
            parent.children.append(item)
        self._write_to_disk(item.uid, title, content)
        return item

    def _build_spine(self, tree):
        for item in tree:
            self.spine.append(item.uid)
            self._build_spine(item.children)

    def _write_spine(self):
        spine_xml = self._template('spine')
        spine_xml.xpath("/*/*[1]/*[1]")[0].text = arrow.utcnow().format(
            "YYYY-MM-DDTHH:mm:ss")
        spine_xml.xpath("/*/*[1]/dc:title", namespaces={
            "dc": "http://purl.org/dc/elements/1.1/"})[0].text = self.title
        self._build_spine(self.root)
        for index, uid in enumerate(self.spine):
            index = '{:04}'.format(index)
            etree.SubElement(spine_xml.xpath("/*/*[2]")[0], "item", **{
                "media-type": "application/xhtml+xml",
                "href": uid + ".xhtml",
                "id": index})
            etree.SubElement(
                spine_xml.xpath("/*/*[3]")[0],
                "itemref", idref=index)
        spine_xml.write(self.dir.name + "/content.opf", xml_declaration=True,
                        encoding="utf-8", pretty_print=True)

    def _write_container(self):
        container_xml = self._template('container')
        meta_inf = Path(self.dir.name).resolve() / 'META-INF'
        meta_inf.mkdir()
        container_xml.write(str(meta_inf / 'container.xml'),
                            xml_declaration=True, encoding="utf-8",
                            pretty_print=True)

    def write(self, filename):
        self._write_spine()
        self._write_container()
        with open(self.dir.name + "/mimetype", "w") as file:
            file.write("application/epub+zip")
        with open(self.dir.name + '/stylesheet.css', 'wb') as file:
            file.write(get_data('pyscp', 'resources/stylesheet.css'))
        with open(self.dir.name + '/cover.png', 'wb') as file:
            file.write(get_data('pyscp', 'resources/cover.png'))
        shutil.make_archive(filename, "zip", self.dir.name)
        shutil.move(filename + ".zip", filename + ".epub")

class HtmlParser:

    def __init__(self, html):
        self.soup = bs4(html).find(id='page-content')

    def _tabviews(self):
        for element in self.soup(class_='yui-navset'):
            element.attrs = {'class': 'tabview'}
            titles = [i.text for i in element.find(class_='yui-nav')('em')]
            element.find(class_='yui-nav').decompose()
            element.div.unwrap()
            for index, tab in enumerate(element('div', recursive=False)):
                tab.attrs = {'class': 'tabview-tab'}
                title = bs4('').new_tag('p', **{'class': 'tab-title'})
                title.string = titles[index]
                tab.insert(0, title)

    def _collapsibles(self):
        for element in self.soup(class_='collapsible-block'):
            element.attrs = {'class': 'collapsible'}
            title = bs4().new_tag('p', **{'class': 'collapsible-title'})
            title.string = element.find(class_='collapsible-block-link').text
            body = element.find(class_='collapsible-block-content')
            element.clear()
            element.append(title)
            for child in list(body.contents):
                element.append(child)

    def _footnotes(self):
        for element in self.soup('sup', class_='footnoteref'):
            element.string = element.a.string
        for element in self.soup(class_='footnote-footer'):
            element.attrs = {'class': 'footnote'}
            element.string = ''.join(element.stripped_strings)

    def _links(self):
        for element in self.soup('a'):
            element.name = 'span'
            element.attrs = {'class': 'link'}

    def _quotes(self):
        for element in self.soup('blockquote'):
            element.name = 'div'
            element.attrs = {'class': 'quote'}

    def parse(self):
        for element in self.soup(class_='page-rate-widget-box'):
            element.decompose()
        self._tabviews()
        self._collapsibles()
        self._footnotes()
        self._links()
        self._quotes()
        return str(self.soup)

###############################################################################
# Page Retrieval Functions
###############################################################################

def main():
    cn = SnapshotConnector(
        'www.scp-wiki.net', '/home/anqxyr/heap/_scp/scp-wiki.2015-03-16.db')
    book = Epub('Testity')
    previous = None
    for url, index in zip(cn.list_pages(), range(50)):
        p = cn(url)
        if index % 5 == 0:
            previous = book.add_page(p.title, p.html)
        else:
            book.add_page(p.title, p.html, parent=previous)
    book.write('test')

if __name__ == "__main__":
    use_default_logging()
    main()
<commit_msg>readd toc generation<commit_after>#!/usr/bin/env python3

###############################################################################
# Module Imports
###############################################################################

import arrow
import itertools
import logging
import shutil

from bs4 import BeautifulSoup as bs4
from collections import namedtuple
from lxml import etree, html
from pkgutil import get_data
from pyscp import SnapshotConnector, use_default_logging
from tempfile import TemporaryDirectory
from pathlib import Path

###############################################################################
# Global Constants And Variables
###############################################################################

log = logging.getLogger('pyscp.ebook')

###############################################################################
# Primary Classes
###############################################################################

class Epub:

    item = namedtuple('EbookItem', 'uid title children')

    def __init__(self, title):
        self.dir = TemporaryDirectory()
        self.title = title
        self.lang = 'en'
        self.root = []
        self.uid = ('{:04}'.format(i) for i in itertools.count(1))

    def _template(self, name):
        name_to_file = dict(page='page.xhtml', spine='content.opf',
                            container='container.xml', toc='toc.ncx')
        return etree.ElementTree(etree.fromstring(
            get_data('pyscp', 'resources/templates/' + name_to_file[name])))

    def _write_to_disk(self, uid, title, content):
        page = self._template('page')
        page.xpath("/*/*[1]/*[1]")[0].text = title
        page.xpath("/*/*[2]")[0].append(html.fromstring(content))
        page.write("{}/{}.xhtml".format(self.dir.name, uid))

    def add_page(self, title, content, parent=None):
        item = self.item(next(self.uid), title, [])
        if not parent:
            self.root.append(item)
        else:
            parent.children.append(item)
        self._write_to_disk(item.uid, title, content)
        return item

    def _flat_tree(self, tree, flat=[]):
        for item in tree:
            flat.append(item.uid)
            self._flat_tree(item.children)
        return flat

    def _write_spine(self):
        spine = self._template('spine')
        spine.xpath("/*/*[1]/*[1]")[0].text = arrow.utcnow().format(
            "YYYY-MM-DDTHH:mm:ss")
        spine.xpath("/*/*[1]/dc:title", namespaces={
            "dc": "http://purl.org/dc/elements/1.1/"})[0].text = self.title
        for index, uid in enumerate(self._flat_tree(self.root)):
            index = '{:04}'.format(index)
            etree.SubElement(spine.xpath("/*/*[2]")[0], "item", **{
                "media-type": "application/xhtml+xml",
                "href": uid + ".xhtml",
                "id": index})
            etree.SubElement(
                spine.xpath("/*/*[3]")[0],
                "itemref", idref=index)
        spine.write(self.dir.name + "/content.opf", xml_declaration=True,
                    encoding="utf-8", pretty_print=True)

    def _write_container(self):
        container = self._template('container')
        meta_inf = Path(self.dir.name).resolve() / 'META-INF'
        meta_inf.mkdir()
        container.write(str(meta_inf / 'container.xml'),
                        xml_declaration=True, encoding="utf-8",
                        pretty_print=True)

    def _add_item_to_toc(self, item, node=None):
        if node is None:
            node = self.toc.xpath("/*/*[3]")[0]
        navpoint = etree.SubElement(node, "navPoint", id=item.uid,
                                    playOrder=str(int(item.uid)))
        navlabel = etree.SubElement(navpoint, "navLabel")
        etree.SubElement(navlabel, "text").text = item.title
        etree.SubElement(navpoint, "content", src="{}.xhtml".format(item.uid))
        for child in item.children:
            self._add_item_to_toc(child, navpoint)

    def _write_toc(self):
        self.toc = self._template('toc')
        self.toc.xpath("/*/*[2]/*")[0].text = self.title
        for item in self.root:
            self._add_item_to_toc(item)
        self.toc.write(
            "{}/toc.ncx".format(self.dir.name), xml_declaration=True,
            encoding="utf-8", pretty_print=True)

    def write(self, filename):
        self._write_spine()
        self._write_container()
        self._write_toc()
        with open(self.dir.name + "/mimetype", "w") as file:
            file.write("application/epub+zip")
        with open(self.dir.name + '/stylesheet.css', 'wb') as file:
            file.write(get_data('pyscp', 'resources/stylesheet.css'))
        with open(self.dir.name + '/cover.png', 'wb') as file:
            file.write(get_data('pyscp', 'resources/cover.png'))
        shutil.make_archive(filename, "zip", self.dir.name)
        shutil.move(filename + ".zip", filename + ".epub")

class HtmlParser:

    def __init__(self, html):
        self.soup = bs4(html).find(id='page-content')

    def _tabviews(self):
        for element in self.soup(class_='yui-navset'):
            element.attrs = {'class': 'tabview'}
            titles = [i.text for i in element.find(class_='yui-nav')('em')]
            element.find(class_='yui-nav').decompose()
            element.div.unwrap()
            for index, tab in enumerate(element('div', recursive=False)):
                tab.attrs = {'class': 'tabview-tab'}
                title = bs4('').new_tag('p', **{'class': 'tab-title'})
                title.string = titles[index]
                tab.insert(0, title)

    def _collapsibles(self):
        for element in self.soup(class_='collapsible-block'):
            element.attrs = {'class': 'collapsible'}
            title = bs4().new_tag('p', **{'class': 'collapsible-title'})
            title.string = element.find(class_='collapsible-block-link').text
            body = element.find(class_='collapsible-block-content')
            element.clear()
            element.append(title)
            for child in list(body.contents):
                element.append(child)

    def _footnotes(self):
        for element in self.soup('sup', class_='footnoteref'):
            element.string = element.a.string
        for element in self.soup(class_='footnote-footer'):
            element.attrs = {'class': 'footnote'}
            element.string = ''.join(element.stripped_strings)

    def _links(self):
        for element in self.soup('a'):
            element.name = 'span'
            element.attrs = {'class': 'link'}

    def _quotes(self):
        for element in self.soup('blockquote'):
            element.name = 'div'
            element.attrs = {'class': 'quote'}

    def parse(self):
        for element in self.soup(class_='page-rate-widget-box'):
            element.decompose()
        self._tabviews()
        self._collapsibles()
        self._footnotes()
        self._links()
        self._quotes()
        return str(self.soup)

###############################################################################
# Page Retrieval Functions
###############################################################################

def main():
    cn = SnapshotConnector(
        'www.scp-wiki.net', '/home/anqxyr/heap/_scp/scp-wiki.2015-03-16.db')
    book = Epub('Testity')
    previous = None
    for url, index in zip(cn.list_pages(), range(50)):
        p = cn(url)
        if index % 5 == 0:
            previous = book.add_page(p.title, p.html)
        else:
            book.add_page(p.title, p.html, parent=previous)
    book.write('test')

if __name__ == "__main__":
    use_default_logging()
    main()
<|endoftext|>


<commit_before>#! /usr/bin/env python

import argparse
from functools import partial

def parser(help_text):
    """Parse out command line arguments"""

    class ArgParser(object):

        def __init__(self, argument_parser):
            self.parser = argument_parser
            self.parse = self.parse_args = self.parser.parse_args

            self.arg = self.parser.add_argument
            self.true = partial(self.arg, action='store_true')
            self.int = partial(self.arg, type=int)

        def strings(self, name, help, many='*'):
            self.arg(name, metavar=name, help=help, type=str, nargs=many)

        def answer(self, answers=None):
            yes_no = {
                'n': ('-n', '--no', 'Answer no'),
                'y': ('-y', '--yes', 'Answer yes'),
            }
            options = yes_no
            options.update(answers)
            [self.true(options[o]) for o in options]

    return ArgParser(
        argparse.ArgumentParser(description=help_text.splitlines()[0]))
<commit_msg>Alias strings to "args"<commit_after>#! /usr/bin/env python

import argparse
from functools import partial

def parser(help_text):
    """Parse out command line arguments"""

    class ArgParser(object):

        def __init__(self, argument_parser):
            self.parser = argument_parser
            self.parse = self.parse_args = self.parser.parse_args

            self.arg = self.parser.add_argument
            self.true = partial(self.arg, action='store_true')
            self.int = partial(self.arg, type=int)

        def strings(self, name, help, many='*'):
            self.arg(name, metavar=name, help=help, type=str, nargs=many)

        args = strings

        def answer(self, answers=None):
            yes_no = {
                'n': ('-n', '--no', 'Answer no'),
                'y': ('-y', '--yes', 'Answer yes'),
            }
            options = yes_no
            options.update(answers)
            [self.true(options[o]) for o in options]

    return ArgParser(
        argparse.ArgumentParser(description=help_text.splitlines()[0]))
<|endoftext|>


<commit_before>#! /usr/bin/env python
import serial
import time

# Serial port
N = "/dev/ttyUSB0"

def ints2str(lst):
    '''
    Taking a list of notes/lengths, convert it to a string
    '''
    s = ""
    for i in lst:
        if i < 0 or i > 255:
            raise Exception
        s = s + str(chr(i))
    return s

# do some initialization magic
s = serial.Serial(N, 57600, timeout=4)
# start code
s.write(ints2str([128]))
# Full mode
s.write(ints2str([132]))

# Drive
#s.write(ints2str([137, 1, 44, 128, 0]))
# wait
#s.write(ints2str([156, 1, 144]))
# Turn
#s.write(ints2str([137, 1, 44, 0, 1]))
#wait
#s.write(ints2str([157, 0, 90]))
s.write(ints2str([152, 13]))
s.write(ints2str([137, 1, 44, 128, 0]))
s.write(ints2str([156, 1, 144]))
s.write(ints2str([137, 0, 0, 0, 0]))

<commit_msg>Update demo.py<commit_after>#! /usr/bin/env python
import serial
import time

# Serial port
N = "/dev/ttyUSB0"

def ints2str(lst):
    '''
    Taking a list of notes/lengths, convert it to a string
    '''
    s = ""
    for i in lst:
        if i < 0 or i > 255:
            raise Exception
        s = s + str(chr(i))
    return s

# do some initialization magic
s = serial.Serial(N, 57600, timeout=4)
# start code
s.write(ints2str([128]))
# Full mode
s.write(ints2str([132]))

# Drive
#s.write(ints2str([137, 1, 44, 128, 0]))
# wait
#s.write(ints2str([156, 1, 144]))
# Turn
#s.write(ints2str([137, 1, 44, 0, 1]))
#wait
#s.write(ints2str([157, 0, 90]))
s.write(ints2str([137, 1, 44, 128, 0]))
s.write(ints2str([156, 1, 144]))
s.write(ints2str([137, 1, 44, 0, 0]))

<|endoftext|>


<commit_before>#!/usr/bin/env python

class LustObject(object):
  # executes the command
  def handle(self, arguments): pass
  def print_help(self): pass

class FactorialCommand(LustObject):
  def handle(self, arguments):
    try: argument = int(arguments[0])
    except (ValueError, IndexError):
      print("fact: could not read integer argument.")
      return

    if argument < 0:
      print("fact: argument has to be non-negative!")
      return

    print(self.__calculate_factorial(argument))

  def print_help(self):
    print(" fact <integer>")
    print("   Calculates the factorial of <integer>.")

  def __calculate_factorial(self, argument):
    # Hmmm...
    result = 0
    for i in range(1, argument+1):
      result *= i
    return result

class QuitCommand(LustObject):
  def handle(self, arguments = None):
    print("Bye!")
    exit()
  def print_help(self):
    print(" quit")
    print("   Quits.")

class HelpCommand(LustObject):
  def __init__(self, commands):
    self.commands = commands
  def handle(self, arguments = None):
    print("List of all commands")
    print("--------------------")
    for command in sorted(self.commands):
      self.commands[command].print_help()
  def print_help(self):
    print(" help")
    print("   Prints help for all commands.")

print("Hello! Welcome to the LARICS Universal Shell Terminal (LUST)!")
print("Enter 'help' for a list of commands. Press Ctrl-D or enter 'quit' to quit.")

# dictionary for storing all commands
commands = { }

commands["fact"] = FactorialCommand()
commands["quit"] = QuitCommand()
# help command needs a reference to the parent dictionary in order to call each
# command's print_help() function
commands["help"] = HelpCommand(commands)

while True:
  # read current line and try to extract command name
  try:
    cmd_line = raw_input(">> ")
  except (EOFError):
    break
  arguments = cmd_line.split()
  try: cmd_name = arguments[0].lower()
  except IndexError: continue

  # look up the appropriate command in commands dictionary
  if cmd_name not in commands:
    print("lust: no such command '{}'.".format(cmd_name))
    continue
  else:
    # command found, pass its handler the rest of the read arguments
    commands[cmd_name].handle(arguments[1:])

print
commands["quit"].handle()
<commit_msg>Fixing bug in factorial<commit_after>#!/usr/bin/env python

class LustObject(object):
  # executes the command
  def handle(self, arguments): pass
  def print_help(self): pass

class FactorialCommand(LustObject):
  def handle(self, arguments):
    try: argument = int(arguments[0])
    except (ValueError, IndexError):
      print("fact: could not read integer argument.")
      return

    if argument < 0:
      print("fact: argument has to be non-negative!")
      return

    print(self.__calculate_factorial(argument))

  def print_help(self):
    print(" fact <integer>")
    print("   Calculates the factorial of <integer>.")

  def __calculate_factorial(self, argument):
    # Hmmm...
    result = 1
    for i in range(1, argument+1):
      result *= i
    return result

class QuitCommand(LustObject):
  def handle(self, arguments = None):
    print("Bye!")
    exit()
  def print_help(self):
    print(" quit")
    print("   Quits.")

class HelpCommand(LustObject):
  def __init__(self, commands):
    self.commands = commands
  def handle(self, arguments = None):
    print("List of all commands")
    print("--------------------")
    for command in sorted(self.commands):
      self.commands[command].print_help()
  def print_help(self):
    print(" help")
    print("   Prints help for all commands.")

print("Hello! Welcome to the LARICS Universal Shell Terminal (LUST)!")
print("Enter 'help' for a list of commands. Press Ctrl-D or enter 'quit' to quit.")

# dictionary for storing all commands
commands = { }

commands["fact"] = FactorialCommand()
commands["quit"] = QuitCommand()
# help command needs a reference to the parent dictionary in order to call each
# command's print_help() function
commands["help"] = HelpCommand(commands)

while True:
  # read current line and try to extract command name
  try:
    cmd_line = raw_input(">> ")
  except (EOFError):
    break
  arguments = cmd_line.split()
  try: cmd_name = arguments[0].lower()
  except IndexError: continue

  # look up the appropriate command in commands dictionary
  if cmd_name not in commands:
    print("lust: no such command '{}'.".format(cmd_name))
    continue
  else:
    # command found, pass its handler the rest of the read arguments
    commands[cmd_name].handle(arguments[1:])

print
commands["quit"].handle()
<|endoftext|>


<commit_before>import json
import logging

from flask import Flask, render_template, request, jsonify
from datetime import datetime
from dateutil.parser import parse
from decimal import Decimal

from pynYNAB.schema.budget import Account, Transaction, Payee

import settings
import ynab_client

app = Flask(__name__, template_folder='../html', static_folder='../static')
app.config['DEBUG'] = settings.flask_debug

if settings.sentry_dsn:
    from raven.contrib.flask import Sentry
    sentry = Sentry(app)

@app.route('/')
def route_index():
    return 'hello world'

@app.route('/webhook', methods=['POST'])
def route_webhook():
    expected_delta = 0
    data = json.loads(request.data.decode('utf8'))
    settings.log.debug('webhook type received %s' % data['type'])

    if data['type'] == 'transaction.created':

        # Sync the account so we get the latest payees
        ynab_client.sync()

        # Does this account exist?
        account = ynab_client.getaccount(settings.ynab_account)
        if account == False:
            return jsonify({'error': 'Account not found'} )

        # Work out the Payee Name
        if data['data'].get('merchant'):
            payee_name = data['data']['merchant']['name']
        else:
            # This is a p2p transaction
            if data['data'].get('counterparty'):
                if data['data']['counterparty'].has_key('name'):
                    payee_name = data['data']['counterparty']['name']
                else:
                    payee_name = data['data']['counterparty']['number']
            elif data['data'].get('metadata', {}).get('is_topup') == 'true':
                payee_name = 'Topup'
            else:
                payee_name = 'Unknown Payee'

        # If we are creating the payee, then we need to increase the delta
        if not ynab_client.payeeexists(payee_name):
            settings.log.debug('payee does not exist, will create %s' % payee_name)
            expected_delta += 1

        # Suggested Tags
        suggested_tags = ''
        if data['data']['merchant'] and data['data']['merchant'].get('metadata', {}).get('suggested_tags'):
            suggested_tags = data['data']['merchant']['metadata']['suggested_tags']

        # Emoji!
        emoji = ''
        if data['data']['merchant'] and data['data']['merchant'].get('emoji'):
            emoji = data['data']['merchant']['emoji']

        # Show the local currency in the notes if this is not in the accounts currency
        local_currency = ''
        if data['data']['local_currency'] != data['data']['currency']:
            local_currency = '(%s %s)' % (data['data']['local_currency'], (abs(data['data']['local_amount']) / 100))

        # Either create or get the payee
        entities_payee_id = ynab_client.getpayee(payee_name).id

        # Create the Transaction
        expected_delta += 1
        settings.log.debug('Creating transaction object')
        transaction = Transaction(
            check_number=data['data']['id'],
            entities_account_id=account.id,
            amount=Decimal(data['data']['amount']) / 100,
            date=parse(data['data']['created']),
            entities_payee_id=entities_payee_id,
            imported_date=datetime.now().date(),
            imported_payee=payee_name,
            memo="%s %s %s" % (emoji, suggested_tags, local_currency),
            source="Imported"
        )

    	if settings.clear_on_import:
            settings.log.debug('Setting transaction as cleared')
     	    transaction.cleared='Cleared'

        settings.log.debug('Duplicate detection')
    	if ynab_client.containsDuplicate(transaction):
            settings.log.debug('skipping due to duplicate transaction')
            return jsonify({'error': 'Skipping due to duplicate transaction'} )
        else:
            settings.log.debug('appending and pushing transaction to YNAB. Delta: %s' % expected_delta)
            ynab_client.client.budget.be_transactions.append(transaction)
            ynab_client.client.push(expected_delta)
            return jsonify(data)
    else:
        settings.log.warning('Unsupported webhook type: %s' % data['type'])
        return jsonify({'error': 'Unsupported webhook type'} )
    return ''

if __name__ == "__main__":
    app.run(host="0.0.0.0")
<commit_msg>Don't store 0 amounts<commit_after>import json
import logging

from flask import Flask, render_template, request, jsonify
from datetime import datetime
from dateutil.parser import parse
from decimal import Decimal

from pynYNAB.schema.budget import Account, Transaction, Payee

import settings
import ynab_client

app = Flask(__name__, template_folder='../html', static_folder='../static')
app.config['DEBUG'] = settings.flask_debug

if settings.sentry_dsn:
    from raven.contrib.flask import Sentry
    sentry = Sentry(app)

@app.route('/')
def route_index():
    return 'hello world'

@app.route('/webhook', methods=['POST'])
def route_webhook():
    expected_delta = 0
    data = json.loads(request.data.decode('utf8'))
    settings.log.debug('webhook type received %s' % data['type'])

    if data['type'] == 'transaction.created':

        # Sync the account so we get the latest payees
        ynab_client.sync()

        if data['data']['amount'] == 0:
            return jsonify({'error': 'Amount is 0'} )

        # Does this account exist?
        account = ynab_client.getaccount(settings.ynab_account)
        if account == False:
            return jsonify({'error': 'Account not found'} )

        # Work out the Payee Name
        if data['data'].get('merchant'):
            payee_name = data['data']['merchant']['name']
        else:
            # This is a p2p transaction
            if data['data'].get('counterparty'):
                if data['data']['counterparty'].has_key('name'):
                    payee_name = data['data']['counterparty']['name']
                else:
                    payee_name = data['data']['counterparty']['number']
            elif data['data'].get('metadata', {}).get('is_topup') == 'true':
                payee_name = 'Topup'
            else:
                payee_name = 'Unknown Payee'

        # If we are creating the payee, then we need to increase the delta
        if not ynab_client.payeeexists(payee_name):
            settings.log.debug('payee does not exist, will create %s' % payee_name)
            expected_delta += 1

        # Suggested Tags
        suggested_tags = ''
        if data['data']['merchant'] and data['data']['merchant'].get('metadata', {}).get('suggested_tags'):
            suggested_tags = data['data']['merchant']['metadata']['suggested_tags']

        # Emoji!
        emoji = ''
        if data['data']['merchant'] and data['data']['merchant'].get('emoji'):
            emoji = data['data']['merchant']['emoji']

        # Show the local currency in the notes if this is not in the accounts currency
        local_currency = ''
        if data['data']['local_currency'] != data['data']['currency']:
            local_currency = '(%s %s)' % (data['data']['local_currency'], (abs(data['data']['local_amount']) / 100))

        # Either create or get the payee
        entities_payee_id = ynab_client.getpayee(payee_name).id

        # Create the Transaction
        expected_delta += 1
        settings.log.debug('Creating transaction object')
        transaction = Transaction(
            check_number=data['data']['id'],
            entities_account_id=account.id,
            amount=Decimal(data['data']['amount']) / 100,
            date=parse(data['data']['created']),
            entities_payee_id=entities_payee_id,
            imported_date=datetime.now().date(),
            imported_payee=payee_name,
            memo="%s %s %s" % (emoji, suggested_tags, local_currency),
            source="Imported"
        )

    	if settings.clear_on_import:
            settings.log.debug('Setting transaction as cleared')
     	    transaction.cleared='Cleared'

        settings.log.debug('Duplicate detection')
    	if ynab_client.containsDuplicate(transaction):
            settings.log.debug('skipping due to duplicate transaction')
            return jsonify({'error': 'Skipping due to duplicate transaction'} )
        else:
            settings.log.debug('appending and pushing transaction to YNAB. Delta: %s' % expected_delta)
            ynab_client.client.budget.be_transactions.append(transaction)
            ynab_client.client.push(expected_delta)
            return jsonify(data)
    else:
        settings.log.warning('Unsupported webhook type: %s' % data['type'])
        return jsonify({'error': 'Unsupported webhook type'} )
    return ''

if __name__ == "__main__":
    app.run(host="0.0.0.0")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8

import json, time
from requests.auth import HTTPDigestAuth
import requests, urllib
import argparse
from elasticsearch import Elasticsearch

from database import *

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--timewindow', default=60, type=int)
    parser.add_argument('-r', '--repeat', default=1, type=int)
    parser.add_argument('-mh', '--moloch', default='localhost')
    parser.add_argument('-u', '--user', default='admin')
    parser.add_argument('-p', '--password', default='admin')
    parser.add_argument('-es','--elasticsearch', nargs='+', required=False, default=['localhost'])
    parser.add_argument('-if', '--influx', default='localhost')
    return parser.parse_args()

ARGS = parse_arguments()

WINDOW=ARGS.timewindow
REPEAT=ARGS.repeat

HOST=ARGS.moloch
PORT=8005
SSL=False
USER=ARGS.user
PASS=ARGS.password

INFLUX_SERVER=ARGS.influx
MEASUREMENT = 'moloch_spi_queries'

USE_ES=False
ES_HOSTS=ARGS.elasticsearch
ES_INDEX='moloch_queries'
ES_TYPE='expression'

# a1 = src_ip
FIELD_MAP = {
    'a1': 'src_ip',
    'p1': 'src_port'
}

QUERIES = [
    {
        'expression': 'protocols != syslog && vlan != 3611 && protocols == http && (http.user-agent == www.owasp.org || http.user-agent == sqlmap || http.user-agent == dirbust || http.user-agent == libwww)',
        'query_tag': 'http_scanner_default_ua',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols != syslog && vlan != 3611 && protocols = http && http.user-agent != mozilla && http.user-agent != "BreakingPoint/1.x (http://bpointsys.com/)" && (http.user-agent == python || http.user-agent == CPython)',
        'query_tag': 'http_python_ua',
        'aggregate_field': 'a1'
    },
    {
        'expression': '(port.dst == 80 || port.dst == 443 || port.dst == 8080 || port.dst == 9200 || port.dst == 8086) && protocols != http && vlan != 3611 && vlan != 3604 && (port.dst != 443 && protocols != tls) && tags != "acked-unseen-segment-dst" && tags != "out-of-order-src"',
        'query_tag': 'not_http_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': '(port.dst != 80 && port.dst != 443 && port.dst != 8080 && port.dst != 9200 && port.dst != 8086) && protocols == http && vlan != 3611',
        'query_tag': 'http_non_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols != syslog && vlan != 3611 && protocols == http && port.dst == 8291',
        'query_tag': 'connection_to_known_beef_server',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'port.dst == 25 && protocols != smtp && vlan != 3611 && vlan != 3611',
        'query_tag': 'not_smtp_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols == smtp && (port != 25 || port != 587) && vlan != 3611',
        'query_tag': 'smtp_non_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'port.dst == 22 && protocols != ssh && vlan != 3611',
        'query_tag': 'not_ssh_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols == ssh && port != 22 && vlan != 3611)',
        'query_tag': 'ssh_non_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols == dns && dns.query.type == AXFR && vlan != 3611',
        'query_tag': 'ssh_non_std_port',
        'aggregate_field': 'a1'
    }
]

def getPeriod(interval=60):
    now = int(time.time())
    then = int(now - interval)
    return then, now

def getQueries():
    queries = []
    es = Elasticsearch(hosts=ES_HOSTS)
    results = es.search(index=ES_INDEX, q='*')
    doc = es.get(index=ES_INDEX, doc_type=ES_TYPE, id=1)['_source']
    for key, value in results['hits'].items():
        if key == 'hits':
            for expr in value:
                queries.append(expr['_source'])
    return queries

def flushInflux(bulk):
    if bulk and isinstance(bulk, list):
        timestamp = bulk[0]['time']
        try:
            db = influxData(host=INFLUX_SERVER)
            client = db.returnClient()
            client.write_points(
                bulk,
                time_precision='s'
                )
            print "sent %s items to influxdb on %s with timestamp %s" % (
                str(len(bulk)),
                db.getHOST(),
                timestamp
                )
        except Exception as e:
            print "FAIL: %s" % (e)
    else:
        print "No items to send"

def queryMoloch(url, startTime, stopTime, query):
    result = []
    field = query['aggregate_field']
    fields = {
        'startTime': int(startTime),
        'stopTime': int(stopTime),
        'expression': query['expression'],
        'field': field,
        'counts': 1
    }
    spi_request = "%s?%s" % (url, urllib.urlencode(fields))
    spi_ret = requests.get(spi_request, auth=HTTPDigestAuth(USER, PASS)).text
    json = {}
    for line in spi_ret.splitlines():
        data = line.split(',')
        json = {
            "time": int(stopTime),
            "measurement": MEASUREMENT,
            "tags": {
                FIELD_MAP[field]: data[0],
                "query": query['query_tag']
            },
            "fields": {
                "count": int(data[1].strip())
            }
        }
        result.append(json)
    return result

def main():

    url = 'http://%s:%s/unique.txt' % (HOST, PORT)

    if USE_ES:
        queries = getQueries()
    else:
        queries = QUERIES

    startTime = time.time()
    for step in range(REPEAT):
        startTime -= WINDOW
        stopTime = startTime
        bulk = []
        for query in queries:
            bulk.extend(queryMoloch(url=url, startTime=startTime, stopTime=stopTime, query=query))
        flushInflux(bulk)

if __name__ == "__main__":
    main()
<commit_msg>Catch zone transfer attps<commit_after>#!/usr/bin/env python
# coding: utf-8

import json, time
from requests.auth import HTTPDigestAuth
import requests, urllib
import argparse
from elasticsearch import Elasticsearch

from database import *

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--timewindow', default=60, type=int)
    parser.add_argument('-r', '--repeat', default=1, type=int)
    parser.add_argument('-mh', '--moloch', default='localhost')
    parser.add_argument('-u', '--user', default='admin')
    parser.add_argument('-p', '--password', default='admin')
    parser.add_argument('-es','--elasticsearch', nargs='+', required=False, default=['localhost'])
    parser.add_argument('-if', '--influx', default='localhost')
    return parser.parse_args()

ARGS = parse_arguments()

WINDOW=ARGS.timewindow
REPEAT=ARGS.repeat

HOST=ARGS.moloch
PORT=8005
SSL=False
USER=ARGS.user
PASS=ARGS.password

INFLUX_SERVER=ARGS.influx
MEASUREMENT = 'moloch_spi_queries'

USE_ES=False
ES_HOSTS=ARGS.elasticsearch
ES_INDEX='moloch_queries'
ES_TYPE='expression'

# a1 = src_ip
FIELD_MAP = {
    'a1': 'src_ip',
    'p1': 'src_port'
}

QUERIES = [
    {
        'expression': 'protocols != syslog && vlan != 3611 && protocols == http && (http.user-agent == www.owasp.org || http.user-agent == sqlmap || http.user-agent == dirbust || http.user-agent == libwww)',
        'query_tag': 'http_scanner_default_ua',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols != syslog && vlan != 3611 && protocols = http && http.user-agent != mozilla && http.user-agent != "BreakingPoint/1.x (http://bpointsys.com/)" && (http.user-agent == python || http.user-agent == CPython)',
        'query_tag': 'http_python_ua',
        'aggregate_field': 'a1'
    },
    {
        'expression': '(port.dst == 80 || port.dst == 443 || port.dst == 8080 || port.dst == 9200 || port.dst == 8086) && protocols != http && vlan != 3611 && vlan != 3604 && (port.dst != 443 && protocols != tls) && tags != "acked-unseen-segment-dst" && tags != "out-of-order-src"',
        'query_tag': 'not_http_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': '(port.dst != 80 && port.dst != 443 && port.dst != 8080 && port.dst != 9200 && port.dst != 8086) && protocols == http && vlan != 3611',
        'query_tag': 'http_non_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols != syslog && vlan != 3611 && protocols == http && port.dst == 8291',
        'query_tag': 'connection_to_known_beef_server',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'port.dst == 25 && protocols != smtp && vlan != 3611 && vlan != 3611',
        'query_tag': 'not_smtp_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols == smtp && (port != 25 || port != 587) && vlan != 3611',
        'query_tag': 'smtp_non_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'port.dst == 22 && protocols != ssh && vlan != 3611',
        'query_tag': 'not_ssh_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols == ssh && port != 22 && vlan != 3611)',
        'query_tag': 'ssh_non_std_port',
        'aggregate_field': 'a1'
    },
    {
        'expression': 'protocols == dns && dns.query.type == AXFR && vlan != 3611',
        'query_tag': 'dns_attempted_zone_transfer',
        'aggregate_field': 'a1'
    }
]

def getPeriod(interval=60):
    now = int(time.time())
    then = int(now - interval)
    return then, now

def getQueries():
    queries = []
    es = Elasticsearch(hosts=ES_HOSTS)
    results = es.search(index=ES_INDEX, q='*')
    doc = es.get(index=ES_INDEX, doc_type=ES_TYPE, id=1)['_source']
    for key, value in results['hits'].items():
        if key == 'hits':
            for expr in value:
                queries.append(expr['_source'])
    return queries

def flushInflux(bulk):
    if bulk and isinstance(bulk, list):
        timestamp = bulk[0]['time']
        try:
            db = influxData(host=INFLUX_SERVER)
            client = db.returnClient()
            client.write_points(
                bulk,
                time_precision='s'
                )
            print "sent %s items to influxdb on %s with timestamp %s" % (
                str(len(bulk)),
                db.getHOST(),
                timestamp
                )
        except Exception as e:
            print "FAIL: %s" % (e)
    else:
        print "No items to send"

def queryMoloch(url, startTime, stopTime, query):
    result = []
    field = query['aggregate_field']
    fields = {
        'startTime': int(startTime),
        'stopTime': int(stopTime),
        'expression': query['expression'],
        'field': field,
        'counts': 1
    }
    spi_request = "%s?%s" % (url, urllib.urlencode(fields))
    spi_ret = requests.get(spi_request, auth=HTTPDigestAuth(USER, PASS)).text
    json = {}
    for line in spi_ret.splitlines():
        data = line.split(',')
        json = {
            "time": int(stopTime),
            "measurement": MEASUREMENT,
            "tags": {
                FIELD_MAP[field]: data[0],
                "query": query['query_tag']
            },
            "fields": {
                "count": int(data[1].strip())
            }
        }
        result.append(json)
    return result

def main():

    url = 'http://%s:%s/unique.txt' % (HOST, PORT)

    if USE_ES:
        queries = getQueries()
    else:
        queries = QUERIES

    startTime = time.time()
    for step in range(REPEAT):
        startTime -= WINDOW
        stopTime = startTime
        bulk = []
        for query in queries:
            bulk.extend(queryMoloch(url=url, startTime=startTime, stopTime=stopTime, query=query))
        flushInflux(bulk)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python

import sys
import time
import pepy
import binascii

from pprint import pprint

p = pepy.parse(sys.argv[1])
ep = p.get_entry_point()
byts = p.get_bytes(ep, 8)
print "Timedatestamp: %s" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(p.timedatestamp))
print "Bytes at 0x%x: %s" % (ep, byts)
for sect in p.get_sections():
    pprint(sect)
<commit_msg>Remove unused import.<commit_after>#!/usr/bin/env python

import sys
import time
import pepy

from pprint import pprint

p = pepy.parse(sys.argv[1])
ep = p.get_entry_point()
byts = p.get_bytes(ep, 8)
print "Timedatestamp: %s" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(p.timedatestamp))
print "Bytes at 0x%x: %s" % (ep, byts)
for sect in p.get_sections():
    pprint(sect)
<|endoftext|>


<commit_before># ----------------------------------------------------------------------------
# Copyright (c) 2016--, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

import os

import click

@click.group(help='Tools for working with QIIME files.')
def tools():
    pass

@tools.command(name='export',
               short_help='Export data from a QIIME Artifact or '
                          'Visualization.',
               help="Export data from a QIIME Artifact or Visualization. "
                    "Exporting extracts the data stored in an Artifact or "
                    "Visualization and will support exporting to multiple "
                    "formats in the future. For now, the data is exported in "
                    "the format it is stored within the Artifact or "
                    "Visualization. Use 'qiime tools extract' to extract the "
                    "Artifact or Visualization's entire archive.")
@click.argument('path', type=click.Path(exists=True, dir_okay=False))
@click.option('--output-dir', required=True,
              type=click.Path(exists=False, dir_okay=True),
              help='Directory where data should be exported to')
def export_data(path, output_dir):
    import qiime.sdk

    result = qiime.sdk.Result.load(path)
    result.export_data(output_dir)

@tools.command(name='import',
               short_help='Import data into a new QIIME Artifact.',
               help="Import data to create a new QIIME Artifact. See "
                    "https://docs.qiime2.org/ for usage examples and details "
                    "on the file types and associated semantic types that can "
                    "be imported.")
@click.option('--type', required=True,
              help='The semantic type of the new artifact.')
@click.option('--input-path', required=True,
              type=click.Path(exists=True, dir_okay=True),
              help='Path to file or directory that should be imported.')
@click.option('--output-path', required=True,
              type=click.Path(exists=False, dir_okay=False),
              help='Path where output artifact should be written.')
@click.option('--source-format', required=False,
              help='The format of the data to be imported. If not provided, '
                   'data must be in the format expected by the semantic type '
                   'provided via --type.')
def import_data(type, input_path, output_path, source_format=None):
    import qiime.sdk

    artifact = qiime.sdk.Artifact.import_data(type, input_path,
                                              view_type=source_format)
    artifact.save(output_path)

@tools.command(short_help='Take a peek at a QIIME Artifact or Visualization.',
               help="Display basic information about a QIIME Artifact or "
                    "Visualization, including its UUID and type.")
@click.argument('path', type=click.Path(exists=True, dir_okay=False))
def peek(path):
    import qiime.sdk

    metadata = qiime.sdk.Result.peek(path)

    click.secho("UUID:        ", fg="green", nl=False)
    click.secho(metadata.uuid)
    click.secho("Type:        ", fg="green", nl=False)
    click.secho(metadata.type)
    if metadata.format is not None:
        click.secho("Data format: ", fg="green", nl=False)
        click.secho(metadata.format)

@tools.command(short_help='View a QIIME Visualization.',
               help="Displays a QIIME Visualization until the command exits. "
                    "To open a QIIME Visualization so it can be used after "
                    "the command exits, use 'qiime extract'.")
@click.argument('visualization-path',
                type=click.Path(exists=True, dir_okay=False))
@click.option('--index-extension', required=False, default='html',
              help='The extension of the index file that should be opened. '
                   '[default: html]')
def view(visualization_path, index_extension):
    import zipfile
    import qiime.sdk

    if index_extension.startswith('.'):
        index_extension = index_extension[1:]
    try:
        visualization = qiime.sdk.Visualization.load(visualization_path)
    # TODO: currently a KeyError is raised if a zipped file that is not a
    # QIIME result is passed. This should be handled better by the framework.
    except (zipfile.BadZipFile, KeyError, TypeError):
        raise click.BadParameter(
            '%s is not a QIIME Visualization. Only QIIME Visualizations can '
            'be viewed.' % visualization_path)

    index_paths = visualization.get_index_paths(relative=False)

    if index_extension not in index_paths:
        raise click.BadParameter(
            'No index %s file with is present in the archive. Available index '
            'extensions are: %s' % (index_extension,
                                    ', '.join(index_paths.keys())))
    else:
        index_path = index_paths[index_extension]
        launch_status = click.launch(index_path)
        if launch_status != 0:
            click.echo('Viewing visualization failed while attempting to '
                       'open %s' % index_path, err=True)
        else:
            while True:
                click.echo(
                    "Press the 'q' key, Control-C, or Control-D to quit. This "
                    "view may no longer be accessible or work correctly after "
                    "quitting.", nl=False)
                # There is currently a bug in click.getchar where translation
                # of Control-C and Control-D into KeyboardInterrupt and
                # EOFError (respectively) does not work on Python 3. The code
                # here should continue to work as expected when the bug is
                # fixed in Click.
                #
                # https://github.com/pallets/click/issues/583
                try:
                    char = click.getchar()
                    click.echo()
                    if char in {'q', '\x03', '\x04'}:
                        break
                except (KeyboardInterrupt, EOFError):
                    break

@tools.command(short_help="Extract a QIIME Artifact or Visualization archive.",
               help="Extract all contents of a QIIME Artifact or "
                    "Visualization's archive, including provenance, metadata, "
                    "and actual data. Use 'qiime tools export' to export only "
                    "the data stored in an Artifact or Visualization, with "
                    "the choice of exporting to different formats.")
@click.argument('path', type=click.Path(exists=True, dir_okay=False))
@click.option('--output-dir', required=False,
              type=click.Path(exists=True, dir_okay=True),
              help='Directory where archive should be extracted to '
                   '[default: current working directory]',
              default=os.getcwd())
def extract(path, output_dir):
    import zipfile
    import qiime.sdk

    try:
        extracted_dir = qiime.sdk.Result.extract(path, output_dir)
    except (zipfile.BadZipFile, ValueError):
        raise click.BadParameter(
            '%s is not a valid QIIME Result. Only QIIME Artifacts and '
            'Visualizations can be extracted.' % path)
    else:
        click.echo('Extracted to %s' % extracted_dir)
<commit_msg>ENH: Disallow viewing on headless environments. (#104)<commit_after># ----------------------------------------------------------------------------
# Copyright (c) 2016--, QIIME 2 development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

import os

import click

@click.group(help='Tools for working with QIIME files.')
def tools():
    pass

@tools.command(name='export',
               short_help='Export data from a QIIME Artifact or '
                          'Visualization.',
               help="Export data from a QIIME Artifact or Visualization. "
                    "Exporting extracts the data stored in an Artifact or "
                    "Visualization and will support exporting to multiple "
                    "formats in the future. For now, the data is exported in "
                    "the format it is stored within the Artifact or "
                    "Visualization. Use 'qiime tools extract' to extract the "
                    "Artifact or Visualization's entire archive.")
@click.argument('path', type=click.Path(exists=True, dir_okay=False))
@click.option('--output-dir', required=True,
              type=click.Path(exists=False, dir_okay=True),
              help='Directory where data should be exported to')
def export_data(path, output_dir):
    import qiime.sdk

    result = qiime.sdk.Result.load(path)
    result.export_data(output_dir)

@tools.command(name='import',
               short_help='Import data into a new QIIME Artifact.',
               help="Import data to create a new QIIME Artifact. See "
                    "https://docs.qiime2.org/ for usage examples and details "
                    "on the file types and associated semantic types that can "
                    "be imported.")
@click.option('--type', required=True,
              help='The semantic type of the new artifact.')
@click.option('--input-path', required=True,
              type=click.Path(exists=True, dir_okay=True),
              help='Path to file or directory that should be imported.')
@click.option('--output-path', required=True,
              type=click.Path(exists=False, dir_okay=False),
              help='Path where output artifact should be written.')
@click.option('--source-format', required=False,
              help='The format of the data to be imported. If not provided, '
                   'data must be in the format expected by the semantic type '
                   'provided via --type.')
def import_data(type, input_path, output_path, source_format=None):
    import qiime.sdk

    artifact = qiime.sdk.Artifact.import_data(type, input_path,
                                              view_type=source_format)
    artifact.save(output_path)

@tools.command(short_help='Take a peek at a QIIME Artifact or Visualization.',
               help="Display basic information about a QIIME Artifact or "
                    "Visualization, including its UUID and type.")
@click.argument('path', type=click.Path(exists=True, dir_okay=False))
def peek(path):
    import qiime.sdk

    metadata = qiime.sdk.Result.peek(path)

    click.secho("UUID:        ", fg="green", nl=False)
    click.secho(metadata.uuid)
    click.secho("Type:        ", fg="green", nl=False)
    click.secho(metadata.type)
    if metadata.format is not None:
        click.secho("Data format: ", fg="green", nl=False)
        click.secho(metadata.format)

@tools.command(short_help='View a QIIME Visualization.',
               help="Displays a QIIME Visualization until the command exits. "
                    "To open a QIIME Visualization so it can be used after "
                    "the command exits, use 'qiime extract'.")
@click.argument('visualization-path',
                type=click.Path(exists=True, dir_okay=False))
@click.option('--index-extension', required=False, default='html',
              help='The extension of the index file that should be opened. '
                   '[default: html]')
def view(visualization_path, index_extension):
    # Guard headless envs from having to import anything
    if not os.getenv("DISPLAY"):
        raise click.UsageError(
            'Visualization viewing is currently not supported in headless '
            'environments. You can view Visualizations (and Artifacts) at '
            'https://view.qiime2.org, or move the Visualization to an '
            'environment with a display and view it with `qiime tools view`.')

    import zipfile
    import qiime.sdk

    if index_extension.startswith('.'):
        index_extension = index_extension[1:]
    try:
        visualization = qiime.sdk.Visualization.load(visualization_path)
    # TODO: currently a KeyError is raised if a zipped file that is not a
    # QIIME result is passed. This should be handled better by the framework.
    except (zipfile.BadZipFile, KeyError, TypeError):
        raise click.BadParameter(
            '%s is not a QIIME Visualization. Only QIIME Visualizations can '
            'be viewed.' % visualization_path)

    index_paths = visualization.get_index_paths(relative=False)

    if index_extension not in index_paths:
        raise click.BadParameter(
            'No index %s file with is present in the archive. Available index '
            'extensions are: %s' % (index_extension,
                                    ', '.join(index_paths.keys())))
    else:
        index_path = index_paths[index_extension]
        launch_status = click.launch(index_path)
        if launch_status != 0:
            click.echo('Viewing visualization failed while attempting to '
                       'open %s' % index_path, err=True)
        else:
            while True:
                click.echo(
                    "Press the 'q' key, Control-C, or Control-D to quit. This "
                    "view may no longer be accessible or work correctly after "
                    "quitting.", nl=False)
                # There is currently a bug in click.getchar where translation
                # of Control-C and Control-D into KeyboardInterrupt and
                # EOFError (respectively) does not work on Python 3. The code
                # here should continue to work as expected when the bug is
                # fixed in Click.
                #
                # https://github.com/pallets/click/issues/583
                try:
                    char = click.getchar()
                    click.echo()
                    if char in {'q', '\x03', '\x04'}:
                        break
                except (KeyboardInterrupt, EOFError):
                    break

@tools.command(short_help="Extract a QIIME Artifact or Visualization archive.",
               help="Extract all contents of a QIIME Artifact or "
                    "Visualization's archive, including provenance, metadata, "
                    "and actual data. Use 'qiime tools export' to export only "
                    "the data stored in an Artifact or Visualization, with "
                    "the choice of exporting to different formats.")
@click.argument('path', type=click.Path(exists=True, dir_okay=False))
@click.option('--output-dir', required=False,
              type=click.Path(exists=True, dir_okay=True),
              help='Directory where archive should be extracted to '
                   '[default: current working directory]',
              default=os.getcwd())
def extract(path, output_dir):
    import zipfile
    import qiime.sdk

    try:
        extracted_dir = qiime.sdk.Result.extract(path, output_dir)
    except (zipfile.BadZipFile, ValueError):
        raise click.BadParameter(
            '%s is not a valid QIIME Result. Only QIIME Artifacts and '
            'Visualizations can be extracted.' % path)
    else:
        click.echo('Extracted to %s' % extracted_dir)
<|endoftext|>