<commit_before>import importlib
import inspect

from antlr4 import CommonTokenStream, InputStream
import six
from stix2patterns.grammars.STIXPatternLexer import STIXPatternLexer
from stix2patterns.grammars.STIXPatternParser import (STIXPatternParser,
                                                      TerminalNode)
from stix2patterns.grammars.STIXPatternVisitor import STIXPatternVisitor
from stix2patterns.validator import STIXPatternErrorListener

from .patterns import (BinaryConstant, BooleanConstant, FloatConstant,
                       FollowedByObservationExpression, HexConstant,
                       IntegerConstant, RepeatQualifier, StartStopQualifier,
                       StringConstant, TimestampConstant, WithinQualifier,
                       _BooleanExpression)

def collapse_lists(lists):
    result = []
    for c in lists:
        if isinstance(c, list):
            result.extend(c)
        else:
            result.append(c)
    return result

def remove_terminal_nodes(parse_tree_nodes):
    values = []
    for x in parse_tree_nodes:
        if not isinstance(x, TerminalNode):
            values.append(x)
    return values

# This class defines a complete generic visitor for a parse tree produced by STIXPatternParser.

class STIXPatternVisitorForSTIX2(STIXPatternVisitor):
    classes = {}

    def __init__(self, module_suffix, module_name):
        if module_suffix and module_name:
            self.module_suffix = module_suffix
            if not STIXPatternVisitorForSTIX2.classes:
                module = importlib.import_module(module_name)
                for k, c in inspect.getmembers(module, inspect.isclass):
                    STIXPatternVisitorForSTIX2.classes[k] = c
        else:
            self.module_suffix = None
        super(STIXPatternVisitor, self).__init__()

    def get_class(self, class_name):
        if class_name in STIXPatternVisitorForSTIX2.classes:
            return STIXPatternVisitorForSTIX2.classes[class_name]
        else:
            return None

    def instantiate(self, klass_name, *args):
        klass_to_instantiate = None
        if self.module_suffix:
            klass_to_instantiate = self.get_class(klass_name + "For" + self.module_suffix)
        if not klass_to_instantiate:
            # use the classes in python_stix2
            klass_to_instantiate = globals()[klass_name]
        return klass_to_instantiate(*args)

    # Visit a parse tree produced by STIXPatternParser#pattern.
    def visitPattern(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    # Visit a parse tree produced by STIXPatternParser#observationExpressions.
    def visitObservationExpressions(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            return FollowedByObservationExpression([children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionOr.
    def visitObservationExpressionOr(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            return self.instantiate("OrObservationExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionAnd.
    def visitObservationExpressionAnd(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            return self.instantiate("AndObservationExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionRepeated.
    def visitObservationExpressionRepeated(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("QualifiedObservationExpression", children[0], children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionSimple.
    def visitObservationExpressionSimple(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ObservationExpression", children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionCompound.
    def visitObservationExpressionCompound(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ObservationExpression", children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionWithin.
    def visitObservationExpressionWithin(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("QualifiedObservationExpression", children[0], children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionStartStop.
    def visitObservationExpressionStartStop(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("QualifiedObservationExpression", children[0], children[1])

    # Visit a parse tree produced by STIXPatternParser#comparisonExpression.
    def visitComparisonExpression(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            if isinstance(children[0], _BooleanExpression):
                children[0].operands.append(children[2])
                return children[0]
            else:
                return self.instantiate("OrBooleanExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#comparisonExpressionAnd.
    def visitComparisonExpressionAnd(self, ctx):
        # TODO: NOT
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            if isinstance(children[0], _BooleanExpression):
                children[0].operands.append(children[2])
                return children[0]
            else:
                return self.instantiate("AndBooleanExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#propTestEqual.
    def visitPropTestEqual(self, ctx):
        children = self.visitChildren(ctx)
        operator = children[1].symbol.type
        negated = operator != STIXPatternParser.EQ
        return self.instantiate("EqualityComparisonExpression", children[0], children[3 if len(children) > 3 else 2],
                                negated)

    # Visit a parse tree produced by STIXPatternParser#propTestOrder.
    def visitPropTestOrder(self, ctx):
        children = self.visitChildren(ctx)
        operator = children[1].symbol.type
        if operator == STIXPatternParser.GT:
            return self.instantiate("GreaterThanComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)
        elif operator == STIXPatternParser.LT:
            return self.instantiate("LessThanComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)
        elif operator == STIXPatternParser.GE:
            return self.instantiate("GreaterThanEqualComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)
        elif operator == STIXPatternParser.LE:
            return self.instantiate("LessThanEqualComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)

    # Visit a parse tree produced by STIXPatternParser#propTestSet.
    def visitPropTestSet(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("InComparisonExpression", children[0], children[3 if len(children) > 3 else 2], False)

    # Visit a parse tree produced by STIXPatternParser#propTestLike.
    def visitPropTestLike(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("LikeComparisonExpression", children[0], children[3 if len(children) > 3 else 2], False)

    # Visit a parse tree produced by STIXPatternParser#propTestRegex.
    def visitPropTestRegex(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("MatchesComparisonExpression", children[0], children[3 if len(children) > 3 else 2],
                                False)

    # Visit a parse tree produced by STIXPatternParser#propTestIsSubset.
    def visitPropTestIsSubset(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("IsSubsetComparisonExpression", children[0], children[3 if len(children) > 3 else 2])

    # Visit a parse tree produced by STIXPatternParser#propTestIsSuperset.
    def visitPropTestIsSuperset(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("IsSupersetComparisonExpression", children[0], children[3 if len(children) > 3 else 2])

    # Visit a parse tree produced by STIXPatternParser#propTestParen.
    def visitPropTestParen(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ParentheticalExpression", children[1])

    # Visit a parse tree produced by STIXPatternParser#startStopQualifier.
    def visitStartStopQualifier(self, ctx):
        children = self.visitChildren(ctx)
        return StartStopQualifier(children[1], children[3])

    # Visit a parse tree produced by STIXPatternParser#withinQualifier.
    def visitWithinQualifier(self, ctx):
        children = self.visitChildren(ctx)
        return WithinQualifier(children[1])

    # Visit a parse tree produced by STIXPatternParser#repeatedQualifier.
    def visitRepeatedQualifier(self, ctx):
        children = self.visitChildren(ctx)
        return RepeatQualifier(children[1])

    # Visit a parse tree produced by STIXPatternParser#objectPath.
    def visitObjectPath(self, ctx):
        children = self.visitChildren(ctx)
        flat_list = collapse_lists(children[2:])
        property_path = []
        i = 0
        while i < len(flat_list):
            current = flat_list[i]
            if i == len(flat_list)-1:
                property_path.append(current)
                break
            next = flat_list[i+1]
            if isinstance(next, TerminalNode):
                property_path.append(self.instantiate("ListObjectPathComponent", current.property_name, next.getText()))
                i += 2
            else:
                property_path.append(current)
                i += 1
        return self.instantiate("ObjectPath", children[0].getText(), property_path)

    # Visit a parse tree produced by STIXPatternParser#objectType.
    def visitObjectType(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    # Visit a parse tree produced by STIXPatternParser#firstPathComponent.
    def visitFirstPathComponent(self, ctx):
        children = self.visitChildren(ctx)
        step = children[0].getText()
        # if step.endswith("_ref"):
        #     return stix2.ReferenceObjectPathComponent(step)
        # else:
        return self.instantiate("BasicObjectPathComponent", step, False)

    # Visit a parse tree produced by STIXPatternParser#indexPathStep.
    def visitIndexPathStep(self, ctx):
        children = self.visitChildren(ctx)
        return children[1]

    # Visit a parse tree produced by STIXPatternParser#pathStep.
    def visitPathStep(self, ctx):
        return collapse_lists(self.visitChildren(ctx))

    # Visit a parse tree produced by STIXPatternParser#keyPathStep.
    def visitKeyPathStep(self, ctx):
        children = self.visitChildren(ctx)
        if isinstance(children[1], StringConstant):
            # special case for hashes
            return children[1].value
        else:
            return self.instantiate("BasicObjectPathComponent", children[1].getText(), True)

    # Visit a parse tree produced by STIXPatternParser#setLiteral.
    def visitSetLiteral(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ListConstant", remove_terminal_nodes(children))

    # Visit a parse tree produced by STIXPatternParser#primitiveLiteral.
    def visitPrimitiveLiteral(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    # Visit a parse tree produced by STIXPatternParser#orderableLiteral.
    def visitOrderableLiteral(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    def visitTerminal(self, node):
        if node.symbol.type == STIXPatternParser.IntPosLiteral or node.symbol.type == STIXPatternParser.IntNegLiteral:
            return IntegerConstant(node.getText())
        elif node.symbol.type == STIXPatternParser.FloatPosLiteral or node.symbol.type == STIXPatternParser.FloatNegLiteral:
            return FloatConstant(node.getText())
        elif node.symbol.type == STIXPatternParser.HexLiteral:
            return HexConstant(node.getText(), from_parse_tree=True)
        elif node.symbol.type == STIXPatternParser.BinaryLiteral:
            return BinaryConstant(node.getText(), from_parse_tree=True)
        elif node.symbol.type == STIXPatternParser.StringLiteral:
            return StringConstant(node.getText().strip('\''), from_parse_tree=True)
        elif node.symbol.type == STIXPatternParser.BoolLiteral:
            return BooleanConstant(node.getText())
        elif node.symbol.type == STIXPatternParser.TimestampLiteral:
            return TimestampConstant(node.getText())
        else:
            return node

    def aggregateResult(self, aggregate, nextResult):
        if aggregate:
            aggregate.append(nextResult)
        elif nextResult:
            aggregate = [nextResult]
        return aggregate

def create_pattern_object(pattern, module_suffix="", module_name=""):
    '''
    Validates a pattern against the STIX Pattern grammar.  Error messages are
    returned in a list.  The test passed if the returned list is empty.
    '''

    start = ''
    if isinstance(pattern, six.string_types):
        start = pattern[:2]
        pattern = InputStream(pattern)

    if not start:
        start = pattern.readline()[:2]
        pattern.seek(0)

    parseErrListener = STIXPatternErrorListener()

    lexer = STIXPatternLexer(pattern)
    # it always adds a console listener by default... remove it.
    lexer.removeErrorListeners()

    stream = CommonTokenStream(lexer)

    parser = STIXPatternParser(stream)
    parser.buildParseTrees = True
    # it always adds a console listener by default... remove it.
    parser.removeErrorListeners()
    parser.addErrorListener(parseErrListener)

    # To improve error messages, replace "<INVALID>" in the literal
    # names with symbolic names.  This is a hack, but seemed like
    # the simplest workaround.
    for i, lit_name in enumerate(parser.literalNames):
        if lit_name == u"<INVALID>":
            parser.literalNames[i] = parser.symbolicNames[i]

    tree = parser.pattern()
    builder = STIXPatternVisitorForSTIX2(module_suffix, module_name)
    return builder.visit(tree)
<commit_msg>Favor star import, disable messages for undefined methods F405<commit_after>import importlib
import inspect

from antlr4 import CommonTokenStream, InputStream
import six
from stix2patterns.grammars.STIXPatternLexer import STIXPatternLexer
from stix2patterns.grammars.STIXPatternParser import (STIXPatternParser,
                                                      TerminalNode)
from stix2patterns.grammars.STIXPatternVisitor import STIXPatternVisitor
from stix2patterns.validator import STIXPatternErrorListener

from .patterns import *
from .patterns import _BooleanExpression

# flake8: noqa F405

def collapse_lists(lists):
    result = []
    for c in lists:
        if isinstance(c, list):
            result.extend(c)
        else:
            result.append(c)
    return result

def remove_terminal_nodes(parse_tree_nodes):
    values = []
    for x in parse_tree_nodes:
        if not isinstance(x, TerminalNode):
            values.append(x)
    return values

# This class defines a complete generic visitor for a parse tree produced by STIXPatternParser.

class STIXPatternVisitorForSTIX2(STIXPatternVisitor):
    classes = {}

    def __init__(self, module_suffix, module_name):
        if module_suffix and module_name:
            self.module_suffix = module_suffix
            if not STIXPatternVisitorForSTIX2.classes:
                module = importlib.import_module(module_name)
                for k, c in inspect.getmembers(module, inspect.isclass):
                    STIXPatternVisitorForSTIX2.classes[k] = c
        else:
            self.module_suffix = None
        super(STIXPatternVisitor, self).__init__()

    def get_class(self, class_name):
        if class_name in STIXPatternVisitorForSTIX2.classes:
            return STIXPatternVisitorForSTIX2.classes[class_name]
        else:
            return None

    def instantiate(self, klass_name, *args):
        klass_to_instantiate = None
        if self.module_suffix:
            klass_to_instantiate = self.get_class(klass_name + "For" + self.module_suffix)
        if not klass_to_instantiate:
            # use the classes in python_stix2
            klass_to_instantiate = globals()[klass_name]
        return klass_to_instantiate(*args)

    # Visit a parse tree produced by STIXPatternParser#pattern.
    def visitPattern(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    # Visit a parse tree produced by STIXPatternParser#observationExpressions.
    def visitObservationExpressions(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            return FollowedByObservationExpression([children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionOr.
    def visitObservationExpressionOr(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            return self.instantiate("OrObservationExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionAnd.
    def visitObservationExpressionAnd(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            return self.instantiate("AndObservationExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionRepeated.
    def visitObservationExpressionRepeated(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("QualifiedObservationExpression", children[0], children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionSimple.
    def visitObservationExpressionSimple(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ObservationExpression", children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionCompound.
    def visitObservationExpressionCompound(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ObservationExpression", children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionWithin.
    def visitObservationExpressionWithin(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("QualifiedObservationExpression", children[0], children[1])

    # Visit a parse tree produced by STIXPatternParser#observationExpressionStartStop.
    def visitObservationExpressionStartStop(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("QualifiedObservationExpression", children[0], children[1])

    # Visit a parse tree produced by STIXPatternParser#comparisonExpression.
    def visitComparisonExpression(self, ctx):
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            if isinstance(children[0], _BooleanExpression):
                children[0].operands.append(children[2])
                return children[0]
            else:
                return self.instantiate("OrBooleanExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#comparisonExpressionAnd.
    def visitComparisonExpressionAnd(self, ctx):
        # TODO: NOT
        children = self.visitChildren(ctx)
        if len(children) == 1:
            return children[0]
        else:
            if isinstance(children[0], _BooleanExpression):
                children[0].operands.append(children[2])
                return children[0]
            else:
                return self.instantiate("AndBooleanExpression", [children[0], children[2]])

    # Visit a parse tree produced by STIXPatternParser#propTestEqual.
    def visitPropTestEqual(self, ctx):
        children = self.visitChildren(ctx)
        operator = children[1].symbol.type
        negated = operator != STIXPatternParser.EQ
        return self.instantiate("EqualityComparisonExpression", children[0], children[3 if len(children) > 3 else 2],
                                negated)

    # Visit a parse tree produced by STIXPatternParser#propTestOrder.
    def visitPropTestOrder(self, ctx):
        children = self.visitChildren(ctx)
        operator = children[1].symbol.type
        if operator == STIXPatternParser.GT:
            return self.instantiate("GreaterThanComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)
        elif operator == STIXPatternParser.LT:
            return self.instantiate("LessThanComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)
        elif operator == STIXPatternParser.GE:
            return self.instantiate("GreaterThanEqualComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)
        elif operator == STIXPatternParser.LE:
            return self.instantiate("LessThanEqualComparisonExpression", children[0],
                                    children[3 if len(children) > 3 else 2], False)

    # Visit a parse tree produced by STIXPatternParser#propTestSet.
    def visitPropTestSet(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("InComparisonExpression", children[0], children[3 if len(children) > 3 else 2], False)

    # Visit a parse tree produced by STIXPatternParser#propTestLike.
    def visitPropTestLike(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("LikeComparisonExpression", children[0], children[3 if len(children) > 3 else 2], False)

    # Visit a parse tree produced by STIXPatternParser#propTestRegex.
    def visitPropTestRegex(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("MatchesComparisonExpression", children[0], children[3 if len(children) > 3 else 2],
                                False)

    # Visit a parse tree produced by STIXPatternParser#propTestIsSubset.
    def visitPropTestIsSubset(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("IsSubsetComparisonExpression", children[0], children[3 if len(children) > 3 else 2])

    # Visit a parse tree produced by STIXPatternParser#propTestIsSuperset.
    def visitPropTestIsSuperset(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("IsSupersetComparisonExpression", children[0], children[3 if len(children) > 3 else 2])

    # Visit a parse tree produced by STIXPatternParser#propTestParen.
    def visitPropTestParen(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ParentheticalExpression", children[1])

    # Visit a parse tree produced by STIXPatternParser#startStopQualifier.
    def visitStartStopQualifier(self, ctx):
        children = self.visitChildren(ctx)
        return StartStopQualifier(children[1], children[3])

    # Visit a parse tree produced by STIXPatternParser#withinQualifier.
    def visitWithinQualifier(self, ctx):
        children = self.visitChildren(ctx)
        return WithinQualifier(children[1])

    # Visit a parse tree produced by STIXPatternParser#repeatedQualifier.
    def visitRepeatedQualifier(self, ctx):
        children = self.visitChildren(ctx)
        return RepeatQualifier(children[1])

    # Visit a parse tree produced by STIXPatternParser#objectPath.
    def visitObjectPath(self, ctx):
        children = self.visitChildren(ctx)
        flat_list = collapse_lists(children[2:])
        property_path = []
        i = 0
        while i < len(flat_list):
            current = flat_list[i]
            if i == len(flat_list)-1:
                property_path.append(current)
                break
            next = flat_list[i+1]
            if isinstance(next, TerminalNode):
                property_path.append(self.instantiate("ListObjectPathComponent", current.property_name, next.getText()))
                i += 2
            else:
                property_path.append(current)
                i += 1
        return self.instantiate("ObjectPath", children[0].getText(), property_path)

    # Visit a parse tree produced by STIXPatternParser#objectType.
    def visitObjectType(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    # Visit a parse tree produced by STIXPatternParser#firstPathComponent.
    def visitFirstPathComponent(self, ctx):
        children = self.visitChildren(ctx)
        step = children[0].getText()
        # if step.endswith("_ref"):
        #     return stix2.ReferenceObjectPathComponent(step)
        # else:
        return self.instantiate("BasicObjectPathComponent", step, False)

    # Visit a parse tree produced by STIXPatternParser#indexPathStep.
    def visitIndexPathStep(self, ctx):
        children = self.visitChildren(ctx)
        return children[1]

    # Visit a parse tree produced by STIXPatternParser#pathStep.
    def visitPathStep(self, ctx):
        return collapse_lists(self.visitChildren(ctx))

    # Visit a parse tree produced by STIXPatternParser#keyPathStep.
    def visitKeyPathStep(self, ctx):
        children = self.visitChildren(ctx)
        if isinstance(children[1], StringConstant):
            # special case for hashes
            return children[1].value
        else:
            return self.instantiate("BasicObjectPathComponent", children[1].getText(), True)

    # Visit a parse tree produced by STIXPatternParser#setLiteral.
    def visitSetLiteral(self, ctx):
        children = self.visitChildren(ctx)
        return self.instantiate("ListConstant", remove_terminal_nodes(children))

    # Visit a parse tree produced by STIXPatternParser#primitiveLiteral.
    def visitPrimitiveLiteral(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    # Visit a parse tree produced by STIXPatternParser#orderableLiteral.
    def visitOrderableLiteral(self, ctx):
        children = self.visitChildren(ctx)
        return children[0]

    def visitTerminal(self, node):
        if node.symbol.type == STIXPatternParser.IntPosLiteral or node.symbol.type == STIXPatternParser.IntNegLiteral:
            return IntegerConstant(node.getText())
        elif node.symbol.type == STIXPatternParser.FloatPosLiteral or node.symbol.type == STIXPatternParser.FloatNegLiteral:
            return FloatConstant(node.getText())
        elif node.symbol.type == STIXPatternParser.HexLiteral:
            return HexConstant(node.getText(), from_parse_tree=True)
        elif node.symbol.type == STIXPatternParser.BinaryLiteral:
            return BinaryConstant(node.getText(), from_parse_tree=True)
        elif node.symbol.type == STIXPatternParser.StringLiteral:
            return StringConstant(node.getText().strip('\''), from_parse_tree=True)
        elif node.symbol.type == STIXPatternParser.BoolLiteral:
            return BooleanConstant(node.getText())
        elif node.symbol.type == STIXPatternParser.TimestampLiteral:
            return TimestampConstant(node.getText())
        else:
            return node

    def aggregateResult(self, aggregate, nextResult):
        if aggregate:
            aggregate.append(nextResult)
        elif nextResult:
            aggregate = [nextResult]
        return aggregate

def create_pattern_object(pattern, module_suffix="", module_name=""):
    '''
    Validates a pattern against the STIX Pattern grammar.  Error messages are
    returned in a list.  The test passed if the returned list is empty.
    '''

    start = ''
    if isinstance(pattern, six.string_types):
        start = pattern[:2]
        pattern = InputStream(pattern)

    if not start:
        start = pattern.readline()[:2]
        pattern.seek(0)

    parseErrListener = STIXPatternErrorListener()

    lexer = STIXPatternLexer(pattern)
    # it always adds a console listener by default... remove it.
    lexer.removeErrorListeners()

    stream = CommonTokenStream(lexer)

    parser = STIXPatternParser(stream)
    parser.buildParseTrees = True
    # it always adds a console listener by default... remove it.
    parser.removeErrorListeners()
    parser.addErrorListener(parseErrListener)

    # To improve error messages, replace "<INVALID>" in the literal
    # names with symbolic names.  This is a hack, but seemed like
    # the simplest workaround.
    for i, lit_name in enumerate(parser.literalNames):
        if lit_name == u"<INVALID>":
            parser.literalNames[i] = parser.symbolicNames[i]

    tree = parser.pattern()
    builder = STIXPatternVisitorForSTIX2(module_suffix, module_name)
    return builder.visit(tree)
<|endoftext|>


<commit_before>""" Bias Baseline
    -------------

    Implementation of bias baseline, in which the prediction is composed by a
    sum of overall average, user bias and item bias, the last two being
    variables to be fitted.

    Usage:
    $ python -m util.bias
"""

from math import sqrt
from sys import argv, exit

from numpy import nan, isnan
from numpy.random import random
from pickle import load

from perf.metrics import calculate_rmse, calculate_ndcg

_ITER = 1000      # number of iterations of stochastic gradient descent
_ALPHA = 0.01     # starting learning rate (MOGHADDAM, with update)
_BETA = 0.01      # regularization factor (MOGHADDAM)
_SAMPLE = 0.001
_TOL = 1e-6
_PKL_DIR = 'out/pkl'
_OUTPUT_DIR = 'out/pred'

class BiasModel(object):
  """ Class implementing a Bias Baseline Model. """

  def __init__(self):
    """ Discriminates existing attributes, initilizing all to None.

        Args:
          None.

        Returns:
          None.
    """
    self.item_bias = None
    self.user_bias = None
    self.overall_mean = None

  def _initialize(self, votes):
    """ Initializes bias values for each entity and calculates overall average.

        Observation:
        - Instead of using random values, static estimation is used by averaging
          deviation from mean (or mean + previous bias).

        Args:
          votes: list of votes (training set).

        Returns:
          None. Instance fields are updated.
    """
    self.user_bias = {}
    user_count = {}
    self.item_bias = {}
    item_count = {}
    self.overall_mean = 0
    count = 0
    for vote in votes:
      self.overall_mean += vote['vote']
      count += 1
    self.overall_mean /= float(count)
    for vote in votes:
      user = vote['voter']
      if user not in self.user_bias:
        self.user_bias[user] = 0
        user_count[user] = 0
      self.user_bias[user] += (vote['vote'] - self.overall_mean)
      user_count[user] += 1
    for user in self.user_bias:
      self.user_bias[user] /= float(user_count[user])
    for vote in votes:
      user = vote['voter']
      item = vote['review']
      if item not in self.item_bias:
        self.item_bias[item] = 0
        item_count[item] = 0
      self.item_bias[item] += (vote['vote'] - self.overall_mean -
          self.user_bias[user])
      item_count[item] += 1
    for item in self.item_bias:
      self.item_bias[item] /= float(item_count[item])

  def fit(self, votes):
    """ Fits a Bias Baseline model given training set (votes).

        Args:
          vote: list of votes, represented as dictionaries (training set).

        Returns:
          None. Instance fields are updated.
    """
    self._initialize(votes)
    previous = float('inf')
    for _ in xrange(_ITER):
      for vote in votes:
        u = vote['voter']
        i = vote['review']
        pred = self.overall_mean + self.user_bias[u] + self.item_bias[i]
        error = float(vote['vote']) - pred
        self.user_bias[u] += _ALPHA * 2 * (error - _BETA * self.user_bias[u])
        self.item_bias[i] += _ALPHA * 2 * (error - _BETA * self.item_bias[i])
      value = 0.0
      for vote in votes:
        u = vote['voter']
        i = vote['review']
        pred = self.overall_mean + self.user_bias[u] + self.item_bias[i]
        value += (vote['vote'] - pred) ** 2
      for u in self.user_bias:
        value += self.user_bias[u] ** 2
      for i in self.item_bias:
        value += self.item_bias[i] ** 2
      if abs(previous - value) < _TOL:
        print 'Convergence'
        break
      previous = value

  def fit_transform(self, votes):
    """ Fits Bias Model on data and adjust vote values to contain only unbiased
        value.

        Args:
          votes: list of votes to learn model and transform values.

        Returns:
          None. Changes are made in place.
    """
    self.fit(votes)
    return self.transform(votes)

  def predict(self, votes):
    """ Predicts a set of vote examples using previous fitted model.

        Args:
          votes: list of dictionaries, representing votes, to predict
        helpfulness vote value.

        Returns:
          A list of floats with predicted vote values.
    """
    pred = []
    for vote in votes:
      u = vote['voter']
      i = vote['review']
      if u in self.user_bias and i in self.item_bias:
        dot = self.overall_mean + self.user_bias[u] + self.item_bias[i]
        pred.append(dot)
      else:
        pred.append(nan)
    return pred

  def get_user_bias(self, user):
    """ Gets a bias regarding a user on a fitted bias model.

        Args:
          user: the id of the user.

        Returns:
          A float with user bias, in warm-start case, or nan, in cold-start.
    """
    return self.user_bias[user] if user in self.user_bias else nan

  def get_item_bias(self, item):
    """ Gets a bias regarding an item on a fitted bias model.

        Args:
          item: the id of the item.

        Returns:
          A float with item bias, in warm-start case, or nan, in cold-start.
    """
    return self.item_bias[item] if item in self.item_bias else nan

  def transform(self, votes):
    """ Removes bias values and overall mean from votes.

        Args:
          votes: list of votes to remove bias; it should be the same list
        fitted, because all entities have to have a bias associated.

        Returns:
          A new list of votes with unbiased values.
    """
    new_votes = []
    for vote in votes:
      new_vote = vote.copy()
      u = vote['voter']
      i = vote['review']
      new_vote['vote'] -= self.overall_mean + self.user_bias[u] + \
          self.item_bias[i]
      new_votes.append(new_vote)
    return new_votes

  def add_bias(self, votes, pred):
    """ Removes bias values and overall mean from votes.

        Args:
          votes: votes whose values are being predicted.
          pred: unbiased estimates of votes.

        Returns:
          None. The pred list is changed in place.
    """
    for index, vote in enumerate(votes):
      u = vote['voter']
      i = vote['review']
      if isnan(pred[index]):
        continue
      pred[index] += self.overall_mean
      if u in self.user_bias:
        pred[index] += self.user_bias[u]
      if i in self.item_bias:
        pred[index] += self.item_bias[i]

if __name__ == '__main__':
  print 'Reading pickles'
  train = load(open('%s/train%.2f.pkl' % (_PKL_DIR, _SAMPLE * 100), 'r'))
  test = load(open('%s/test%.2f.pkl' % (_PKL_DIR, _SAMPLE * 100), 'r'))
  overall_mean = float(sum([float(v['vote']) for v in train])) / len(train)

  print 'Fitting Model'
  model = BiasModel()
  model.fit(train)

  print 'Calculating Predictions'
  pred = model.predict(train)

  print 'TRAINING ERROR'
  truth = [v['vote'] for v in train]
  rmse = calculate_rmse(pred, truth)
  print 'RMSE: %s' % rmse
  for i in xrange(5, 21, 5):
    score = calculate_ndcg(pred, truth, i)
    print 'NDCG@%d: %f' % (i, score)

  print 'Outputing Prediction'
  pred = model.predict(test)
  output = open('%s/bias%.2f.dat' % (_OUTPUT_DIR, _SAMPLE * 100), 'w')
  for p in pred:
    print >> output, overall_mean if isnan(p) else p
  output.close()

<commit_msg>New bias version, replacing review by author and product<commit_after>""" Bias Baseline
    -------------

    Implementation of bias baseline, in which the prediction is composed by a
    sum of overall average, user bias and item bias, the last two being
    variables to be fitted.

    Usage:
    $ python -m util.bias
"""

from math import sqrt
from sys import argv, exit

from numpy import nan, isnan
from numpy.random import random
from pickle import load

from perf.metrics import calculate_rmse, calculate_ndcg

_ITER = 1000      # number of iterations of stochastic gradient descent
_ALPHA = 0.01     # starting learning rate (MOGHADDAM, with update)
_BETA = 0.01      # regularization factor (MOGHADDAM)
_SAMPLE = 0.001
_TOL = 1e-6
_PKL_DIR = 'out/pkl'
_OUTPUT_DIR = 'out/pred'

class BiasModel(object):
  """ Class implementing a Bias Baseline Model. """

  def __init__(self):
    """ Discriminates existing attributes, initilizing all to None.

        Args:
          None.

        Returns:
          None.
    """
    self.product_bias = None
    self.author_bias = None
    self.voter_bias = None
    self.overall_mean = None

  def _initialize(self, votes, reviews):
    """ Initializes bias values for each entity and calculates overall average.

        Observation:
        - Instead of using random values, static estimation is used by averaging
          deviation from mean (or mean + previous bias).

        Args:
          votes: list of votes (training set).
          reviews: dictionary of reviews, to obtain product.

        Returns:
          None. Instance fields are updated.
    """
    self.voter_bias = {}
    voter_count = {}
    self.author_bias = {}
    author_count = {}
    self.product_bias = {}
    product_count = {}
    self.overall_mean = 0
    count = 0
    for vote in votes:
      self.overall_mean += vote['vote']
      count += 1
    self.overall_mean /= float(count)
    for vote in votes:
      voter = vote['voter']
      if voter not in self.voter_bias:
        self.voter_bias[voter] = 0
        voter_count[voter] = 0
      self.voter_bias[voter] += (vote['vote'] - self.overall_mean)
      voter_count[voter] += 1
    for voter in self.voter_bias:
      self.voter_bias[voter] /= float(voter_count[voter])
    for vote in votes:
      voter = vote['voter']
      author = vote['author']
      if author not in self.author_bias:
        self.author_bias[author] = 0
        author_count[author] = 0
      self.author_bias[author] += (vote['vote'] - self.overall_mean -
          self.voter_bias[voter])
      author_count[author] += 1
    for author in self.author_bias:
      self.author_bias[author] /= float(author_count[author])
    for vote in votes:
      voter = vote['voter']
      author = vote['author']
      product = reviews[vote['review']]['product']
      if product not in self.product_bias:
        self.product_bias[product] = 0
        product_count[product] = 0
      self.product_bias[product] += (vote['vote'] - self.overall_mean -
          self.voter_bias[voter] - self.author_bias[author])
      product_count[product] += 1
    for product in self.product_bias:
      self.product_bias[product] /= float(product_count[product])

  def fit(self, votes, reviews):
    """ Fits a Bias Baseline model given training set (votes).

        Args:
          vote: list of votes, represented as dictionaries (training set).
          reviews: dictionary of reviews.

        Returns:
          None. Instance fields are updated.
    """
    self._initialize(votes, reviews)
    previous = float('inf')
    for _ in xrange(_ITER):
      for vote in votes:
        voter = vote['voter']
        author = vote['author']
        product = reviews[vote['review']]['product']
        pred = self.overall_mean + self.voter_bias[voter] + \
            self.author_bias[author] + self.product_bias[product]
        error = pred - float(vote['vote'])
        # one does not depend on the other: no need for temporary variables
        self.voter_bias[voter] -= _ALPHA * 2 * (error + _BETA *
            self.voter_bias[voter])
        self.author_bias[author] -= _ALPHA * 2 * (error + _BETA * \
            self.author_bias[author])
        self.product_bias[product] -= _ALPHA * 2 * (error + _BETA * \
            self.product_bias[product])
      value = 0.0
      for vote in votes:
        voter = vote['voter']
        author = vote['author']
        product = reviews[vote['review']]['product']
        pred = self.overall_mean + self.voter_bias[voter] + \
            self.author_bias[author] + self.product_bias[product]
        value += (vote['vote'] - pred) ** 2
      for voter in self.voter_bias:
        value += self.voter_bias[voter] ** 2
      for author in self.author_bias:
        value += self.author_bias[author] ** 2
      for product in self.product_bias:
        value += self.product_bias[product] ** 2
      if abs(previous - value) < _TOL:
        print 'Convergence'
        break
      previous = value

  def fit_transform(self, votes, reviews):
    """ Fits Bias Model on data and adjust vote values to contain only unbiased
        value.

        Args:
          votes: list of votes to learn model and transform values.
          reviews: dictionary of reviews.

        Returns:
          None. Changes are made in place.
    """
    self.fit(votes, reviews)
    return self.transform(votes, reviews)

  def predict(self, votes, reviews):
    """ Predicts a set of vote examples using previous fitted model.

        Args:
          votes: list of dictionaries, representing votes, to predict
        helpfulness vote value.
          reviews: dictionary of reviews.

        Returns:
          A list of floats with predicted vote values.
    """
    pred = []
    for vote in votes:
      voter = vote['voter']
      author = vote['author']
      product = reviews[vote['review']]['product']
      value = self.overall_mean
      if voter in self.voter_bias: # if no bias, we assume 0
        dot += self.voter_bias[voter]
      if author in self.author_bias:
        dot += self.author_bias[author]
      if product in self.product_bias:
        dot += self.product_bias[product]
      pred.append(value)
    return pred

  def transform(self, votes, reviews):
    """ Removes bias values and overall mean from votes.

        Args:
          votes: list of votes to remove bias; it should be the same list
        fitted, because all entities have to have a bias associated.
          reviews: dictionary of reviews.

        Returns:
          A new list of votes with unbiased values.
    """
    new_votes = []
    for vote in votes:
      new_vote = vote.copy()
      voter = vote['voter']
      author = vote['author']
      product = reviews[vote['review']]['product']
      new_vote['vote'] -= self.overall_mean + self.voter_bias[voter] + \
          self.author_bias[author] + self.product_bias[product]
      new_votes.append(new_vote)
    return new_votes

  def add_bias(self, votes, reviews, pred):
    """ Removes bias values and overall mean from votes.

        Args:
          votes: votes whose values are being predicted.
          reviews: dictionary of reviews.
          pred: unbiased estimates of votes.

        Returns:
          None. The pred list is changed in place.
    """
    for index, vote in enumerate(votes):
      voter = vote['voter']
      author = vote['author']
      product = reviews[vote['review']]['product']
      if isnan(pred[index]):
        continue
      pred[index] += self.overall_mean
      if voter in self.voter_bias:
        pred[index] += self.voter_bias[voter]
      if author in self.author_bias:
        pred[index] += self.author_bias[author]
      if product in self.product_bias:
        pred[index] += self.product_bias[product]

if __name__ == '__main__':
  print 'Reading pickles'
  train = load(open('%s/train%.2f.pkl' % (_PKL_DIR, _SAMPLE * 100), 'r'))
  test = load(open('%s/test%.2f.pkl' % (_PKL_DIR, _SAMPLE * 100), 'r'))
  overall_mean = float(sum([float(v['vote']) for v in train])) / len(train)

  print 'Fitting Model'
  model = BiasModel()
  model.fit(train)

  print 'Calculating Predictions'
  pred = model.predict(train)

  print 'TRAINING ERROR'
  truth = [v['vote'] for v in train]
  rmse = calculate_rmse(pred, truth)
  print 'RMSE: %s' % rmse
  for i in xrange(5, 21, 5):
    score = calculate_ndcg(pred, truth, i)
    print 'NDCG@%d: %f' % (i, score)

  print 'Outputing Prediction'
  pred = model.predict(test)
  output = open('%s/bias%.2f.dat' % (_OUTPUT_DIR, _SAMPLE * 100), 'w')
  for p in pred:
    print >> output, overall_mean if isnan(p) else p
  output.close()

<|endoftext|>


<commit_before># SFTP storage backend for Django.
# Author: Brent Tubbs <brent.tubbs@gmail.com>
# License: MIT
#
# Modeled on the FTP storage by Rafal Jonca <jonca.rafal@gmail.com>
#
# Settings:
#
# SFTP_STORAGE_HOST - The hostname where you want the files to be saved.
#
# SFTP_STORAGE_ROOT - The root directory on the remote host into which files
# should be placed.  Should work the same way that STATIC_ROOT works for local
# files.  Must include a trailing slash.
#
# SFTP_STORAGE_PARAMS (Optional) - A dictionary containing connection
# parameters to be passed as keyword arguments to
# paramiko.SSHClient().connect() (do not include hostname here).  See
# http://www.lag.net/paramiko/docs/paramiko.SSHClient-class.html#connect for
# details
#
# SFTP_STORAGE_INTERACTIVE (Optional) - A boolean indicating whether to prompt
# for a password if the connection cannot be made using keys, and there is not
# already a password in SFTP_STORAGE_PARAMS.  You can set this to True to
# enable interactive login when running 'manage.py collectstatic', for example.
#
#   DO NOT set SFTP_STORAGE_INTERACTIVE to True if you are using this storage
#   for files being uploaded to your site by users, because you'll have no way
#   to enter the password when they submit the form..
#
# SFTP_STORAGE_FILE_MODE (Optional) - A bitmask for setting permissions on
# newly-created files.  See http://docs.python.org/library/os.html#os.chmod for
# acceptable values.
#
# SFTP_STORAGE_DIR_MODE (Optional) - A bitmask for setting permissions on
# newly-created directories.  See
# http://docs.python.org/library/os.html#os.chmod for acceptable values.
#
#   Hint: if you start the mode number with a 0 you can express it in octal
#   just like you would when doing "chmod 775 myfile" from bash.
#
# SFTP_STORAGE_UID (Optional) - uid of the account that should be set as owner
# of the files on the remote host.  You have to be root to set this.
#
# SFTP_STORAGE_GID (Optional) - gid of the group that should be set on the
# files on the remote host.  You have to be a member of the group to set this.

import os
import stat
import paramiko
import posixpath
import getpass
from datetime import datetime
from django.conf import settings
from django.core.files.storage import Storage
from django.core.files.base import File
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

class SFTPStorage(Storage):

    def __init__(self):
        self._host = settings.SFTP_STORAGE_HOST

        # if present, settings.SFTP_STORAGE_PARAMS should be a dict with params
        # matching the keyword arguments to paramiko.SSHClient().connect().  So
        # you can put username/password there.  Or you can omit all that if
        # you're using keys.
        self._params = getattr(settings, 'SFTP_STORAGE_PARAMS', {})
        self._interactive = getattr(settings, 'SFTP_STORAGE_INTERACTIVE',
                                    False)

        self._file_mode = getattr(settings, 'SFTP_STORAGE_FILE_MODE', None)
        self._dir_mode = getattr(settings, 'SFTP_STORAGE_DIR_MODE', None)

        self._uid = getattr(settings, 'SFTP_STORAGE_UID', None)
        self._gid = getattr(settings, 'SFTP_STORAGE_GID', None)

        self._root_path = settings.SFTP_STORAGE_ROOT

        # for now it's all posix paths.  Maybe someday we'll support figuring
        # out if the remote host is windows.
        self._pathmod = posixpath

        # set up connection
        self._connect()

    def _connect(self):
        self._ssh = paramiko.SSHClient()

        # automatically add host keys from current user.
        self._ssh.load_host_keys(os.path.expanduser(os.path.join("~", ".ssh", "known_hosts")))

        # and automatically add new host keys for hosts we haven't seen before.
        self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        try:
            self._ssh.connect(self._host, **self._params)
        except paramiko.AuthenticationException, e:
            if self._interactive and 'password' not in self._params:
                # If authentication has failed, and we haven't already tried
                # username/password, and configuration allows it, then try
                # again with username/password.
                if 'username' not in self._params:
                    self._params['username'] = getpass.getuser()
                self._params['password'] = getpass.getpass()
                self._connect()
            else:
                raise paramiko.AuthenticationException, e

        if not hasattr(self, '_sftp'):
            self._sftp = self._ssh.open_sftp()

    def _join(self, *args):
        # Use the path module for the remote host type to join a path together
        return self._pathmod.join(*args)

    def _remote_path(self, name):
        return self._join(self._root_path, name)

    def _open(self, name, mode='rb'):
        return SFTPStorageFile(name, self, mode)

    def _read(self, name):
        remote_path = self._remote_path(name)
        return self._sftp.open(remote_path, 'rb')

    def _chown(self, path, uid=None, gid=None):
        """Set uid and/or gid for file at path."""
        # Paramiko's chown requires both uid and gid, so look them up first if
        # we're only supposed to set one.
        if uid is None or gid is None:
            attr = self._sftp.stat(path)
            uid = uid or attr.st_uid
            gid = gid or attr.st_gid
        self._sftp.chown(path, uid, gid)

    def _mkdir(self, path):
        """Create directory, recursing up to create parent dirs if
        necessary."""
        parent = self._pathmod.dirname(path)
        if not self.exists(parent):
            self._mkdir(parent)
        self._sftp.mkdir(path)

        if self._dir_mode is not None:
            self._sftp.chmod(path, self._dir_mode)

        if self._uid or self._gid:
            self._chown(path, uid=self._uid, gid=self._gid)

    def _save(self, name, content):
        """Save file via SFTP."""
        content.open()
        path = self._remote_path(name)
        dirname = self._pathmod.dirname(path)
        if not self.exists(dirname):
            self._mkdir(dirname)

        f = self._sftp.open(path, 'wb')
        f.write(content.file.read())
        f.close()

        # set file permissions if configured
        if self._file_mode is not None:
            self._sftp.chmod(path, self._file_mode)
        if self._uid or self._gid:
            self._chown(path, uid=self._uid, gid=self._gid)
        return name

    def delete(self, name):
        remote_path = self._remote_path(name)
        self._sftp.remove(remote_path)

    def exists(self, name):
        # Try to retrieve file info.  Return true on success, false on failure.
        remote_path = self._remote_path(name)
        try:
            self._sftp.stat(remote_path)
            return True
        except IOError:
            return False

    def _isdir_attr(self, item):
        # Return whether an item in sftp.listdir_attr results is a directory
        if item.st_mode is not None:
            return stat.S_IFMT(item.st_mode) == stat.S_IFDIR
        else:
            return False

    def listdir(self, path):
        remote_path = self._remote_path(path)
        dirs, files = [], []
        for item in self._sftp.listdir_attr(remote_path):
            if self._isdir_attr(item):
                dirs.append(item.filename)
            else:
                files.append(item.filename)
        return dirs, files

    def size(self, name):
        remote_path = self._remote_path(name)
        return self._sftp.stat(remote_path).st_size

    def accessed_time(self, name):
        remote_path = self._remote_path(name)
        utime = self._sftp.stat(remote_path).st_atime
        return datetime.fromtimestamp(utime)

    def modified_time(self, name):
        remote_path = self._remote_path(name)
        utime = self._sftp.stat(remote_path).st_mtime
        return datetime.fromtimestamp(utime)

    def url(self, name):
        remote_path = self._remote_path(name)
        return 'sftp://%s/%s' % (self._host, remote_path)

class SFTPStorageFile(File):
    def __init__(self, name, storage, mode):
        self._name = name
        self._storage = storage
        self._mode = mode
        self._is_dirty = False
        self.file = StringIO()
        self._is_read = False

    @property
    def size(self):
        if not hasattr(self, '_size'):
            self._size = self._storage.size(self._name)
        return self._size

    def read(self, num_bytes=None):
        if not self._is_read:
            self.file = self._storage._read(self._name)
            self._is_read = True

        return self.file.read(num_bytes)

    def write(self, content):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self.file = StringIO(content)
        self._is_dirty = True
        self._is_read = True

    def close(self):
        if self._is_dirty:
            self._storage._save(self._name, self.file.getvalue())
        self.file.close()
<commit_msg>use lazy SFTP connection to avoid paramiko hanging on import when storage is passed into file field init<commit_after># SFTP storage backend for Django.
# Author: Brent Tubbs <brent.tubbs@gmail.com>
# License: MIT
#
# Modeled on the FTP storage by Rafal Jonca <jonca.rafal@gmail.com>
#
# Settings:
#
# SFTP_STORAGE_HOST - The hostname where you want the files to be saved.
#
# SFTP_STORAGE_ROOT - The root directory on the remote host into which files
# should be placed.  Should work the same way that STATIC_ROOT works for local
# files.  Must include a trailing slash.
#
# SFTP_STORAGE_PARAMS (Optional) - A dictionary containing connection
# parameters to be passed as keyword arguments to
# paramiko.SSHClient().connect() (do not include hostname here).  See
# http://www.lag.net/paramiko/docs/paramiko.SSHClient-class.html#connect for
# details
#
# SFTP_STORAGE_INTERACTIVE (Optional) - A boolean indicating whether to prompt
# for a password if the connection cannot be made using keys, and there is not
# already a password in SFTP_STORAGE_PARAMS.  You can set this to True to
# enable interactive login when running 'manage.py collectstatic', for example.
#
#   DO NOT set SFTP_STORAGE_INTERACTIVE to True if you are using this storage
#   for files being uploaded to your site by users, because you'll have no way
#   to enter the password when they submit the form..
#
# SFTP_STORAGE_FILE_MODE (Optional) - A bitmask for setting permissions on
# newly-created files.  See http://docs.python.org/library/os.html#os.chmod for
# acceptable values.
#
# SFTP_STORAGE_DIR_MODE (Optional) - A bitmask for setting permissions on
# newly-created directories.  See
# http://docs.python.org/library/os.html#os.chmod for acceptable values.
#
#   Hint: if you start the mode number with a 0 you can express it in octal
#   just like you would when doing "chmod 775 myfile" from bash.
#
# SFTP_STORAGE_UID (Optional) - uid of the account that should be set as owner
# of the files on the remote host.  You have to be root to set this.
#
# SFTP_STORAGE_GID (Optional) - gid of the group that should be set on the
# files on the remote host.  You have to be a member of the group to set this.

import os
import stat
import paramiko
import posixpath
import getpass
from datetime import datetime
from django.conf import settings
from django.core.files.storage import Storage
from django.core.files.base import File
try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

class SFTPStorage(Storage):

    def __init__(self):
        self._host = settings.SFTP_STORAGE_HOST

        # if present, settings.SFTP_STORAGE_PARAMS should be a dict with params
        # matching the keyword arguments to paramiko.SSHClient().connect().  So
        # you can put username/password there.  Or you can omit all that if
        # you're using keys.
        self._params = getattr(settings, 'SFTP_STORAGE_PARAMS', {})
        self._interactive = getattr(settings, 'SFTP_STORAGE_INTERACTIVE',
                                    False)
        self._file_mode = getattr(settings, 'SFTP_STORAGE_FILE_MODE', None)
        self._dir_mode = getattr(settings, 'SFTP_STORAGE_DIR_MODE', None)

        self._uid = getattr(settings, 'SFTP_STORAGE_UID', None)
        self._gid = getattr(settings, 'SFTP_STORAGE_GID', None)

        self._root_path = settings.SFTP_STORAGE_ROOT

        # for now it's all posix paths.  Maybe someday we'll support figuring
        # out if the remote host is windows.
        self._pathmod = posixpath

    def _connect(self):
        self._ssh = paramiko.SSHClient()

        # automatically add host keys from current user.
        self._ssh.load_host_keys(os.path.expanduser(os.path.join("~", ".ssh", "known_hosts")))

        # and automatically add new host keys for hosts we haven't seen before.
        self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        try:
            self._ssh.connect(self._host, **self._params)
        except paramiko.AuthenticationException, e:
            if self._interactive and 'password' not in self._params:
                # If authentication has failed, and we haven't already tried
                # username/password, and configuration allows it, then try
                # again with username/password.
                if 'username' not in self._params:
                    self._params['username'] = getpass.getuser()
                self._params['password'] = getpass.getpass()
                self._connect()
            else:
                raise paramiko.AuthenticationException, e
        except Exception, e:
            print e

        if not hasattr(self, '_sftp'):
            self._sftp = self._ssh.open_sftp()

    @property
    def sftp(self):
        """Lazy SFTP connection"""
        if not hasattr(self, '_sftp'):
            self._connect()
        return self._sftp

    def _join(self, *args):
        # Use the path module for the remote host type to join a path together
        return self._pathmod.join(*args)

    def _remote_path(self, name):
        return self._join(self._root_path, name)

    def _open(self, name, mode='rb'):
        return SFTPStorageFile(name, self, mode)

    def _read(self, name):
        remote_path = self._remote_path(name)
        return self.sftp.open(remote_path, 'rb')

    def _chown(self, path, uid=None, gid=None):
        """Set uid and/or gid for file at path."""
        # Paramiko's chown requires both uid and gid, so look them up first if
        # we're only supposed to set one.
        if uid is None or gid is None:
            attr = self.sftp.stat(path)
            uid = uid or attr.st_uid
            gid = gid or attr.st_gid
        self.sftp.chown(path, uid, gid)

    def _mkdir(self, path):
        """Create directory, recursing up to create parent dirs if
        necessary."""
        parent = self._pathmod.dirname(path)
        if not self.exists(parent):
            self._mkdir(parent)
        self.sftp.mkdir(path)

        if self._dir_mode is not None:
            self.sftp.chmod(path, self._dir_mode)

        if self._uid or self._gid:
            self._chown(path, uid=self._uid, gid=self._gid)

    def _save(self, name, content):
        """Save file via SFTP."""
        content.open()
        path = self._remote_path(name)
        dirname = self._pathmod.dirname(path)
        if not self.exists(dirname):
            self._mkdir(dirname)

        f = self.sftp.open(path, 'wb')
        f.write(content.file.read())
        f.close()

        # set file permissions if configured
        if self._file_mode is not None:
            self.sftp.chmod(path, self._file_mode)
        if self._uid or self._gid:
            self._chown(path, uid=self._uid, gid=self._gid)
        return name

    def delete(self, name):
        remote_path = self._remote_path(name)
        self.sftp.remove(remote_path)

    def exists(self, name):
        # Try to retrieve file info.  Return true on success, false on failure.
        remote_path = self._remote_path(name)
        try:
            self.sftp.stat(remote_path)
            return True
        except IOError:
            return False

    def _isdir_attr(self, item):
        # Return whether an item in sftp.listdir_attr results is a directory
        if item.st_mode is not None:
            return stat.S_IFMT(item.st_mode) == stat.S_IFDIR
        else:
            return False

    def listdir(self, path):
        remote_path = self._remote_path(path)
        dirs, files = [], []
        for item in self.sftp.listdir_attr(remote_path):
            if self._isdir_attr(item):
                dirs.append(item.filename)
            else:
                files.append(item.filename)
        return dirs, files

    def size(self, name):
        remote_path = self._remote_path(name)
        return self.sftp.stat(remote_path).st_size

    def accessed_time(self, name):
        remote_path = self._remote_path(name)
        utime = self.sftp.stat(remote_path).st_atime
        return datetime.fromtimestamp(utime)

    def modified_time(self, name):
        remote_path = self._remote_path(name)
        utime = self.sftp.stat(remote_path).st_mtime
        return datetime.fromtimestamp(utime)

    def url(self, name):
        remote_path = self._remote_path(name)
        return 'sftp://%s/%s' % (self._host, remote_path)

class SFTPStorageFile(File):
    def __init__(self, name, storage, mode):
        self._name = name
        self._storage = storage
        self._mode = mode
        self._is_dirty = False
        self.file = StringIO()
        self._is_read = False

    @property
    def size(self):
        if not hasattr(self, '_size'):
            self._size = self._storage.size(self._name)
        return self._size

    def read(self, num_bytes=None):
        if not self._is_read:
            self.file = self._storage._read(self._name)
            self._is_read = True

        return self.file.read(num_bytes)

    def write(self, content):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self.file = StringIO(content)
        self._is_dirty = True
        self._is_read = True

    def close(self):
        if self._is_dirty:
            self._storage._save(self._name, self.file.getvalue())
        self.file.close()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.contrib import admin
from .models import PageStatistic

class PageStatisticAdmin(admin.ModelAdmin):
    search_fields = ("page",)
    list_display = ("page", "likes", "shares", "parse_date",)

admin.site.register(PageStatistic, PageStatisticAdmin)
<commit_msg>admin fix<commit_after># -*- coding: utf-8 -*-
from django.contrib import admin
from .models import PageStatistic

class PageStatisticAdmin(admin.ModelAdmin):
    search_fields = ("page",)
    list_display = ("page", "likes", "talking_about_count", "updated_at",)

admin.site.register(PageStatistic, PageStatisticAdmin)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2016 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Basic Turbinia config."""

from __future__ import unicode_literal

import imp
import itertools
import logging
import os
import sys

log = logging.getLogger('turbinia')

# Look for config files with these names
CONFIGFILES = ['.turbiniarc', 'turbinia.conf', 'turbinia_config.py']
# Look in homedir first, then /etc/turbinia, and finally in the source
# config dir for config files
CONFIGPATH = [
    os.path.expanduser('~'),
    '/etc/turbinia',
    os.path.dirname(os.path.abspath(__file__))]
# Config vars that we expect to exist in the configuration
CONFIGVARS = [
    # Turbinia Config
    'TASK_MANAGER',
    'LOG_FILE',
    'OUTPUT_DIR',
    'SLEEP_TIME',
    'SINGLE_RUN',
    'MOUNT_DIR_PREFIX',
    'SHARED_FILESYSTEM',
    # GCE CONFIG
    'PROJECT',
    'ZONE',
    'INSTANCE',
    'DEVICE_NAME',
    'SCRATCH_PATH',
    'BUCKET_NAME',
    'PSQ_TOPIC',
    'PUBSUB_TOPIC',
    'GCS_OUTPUT_PATH',
    'STATE_MANAGER',
    # REDIS CONFIG
    'REDIS_HOST',
    'REDIS_PORT',
    # Timesketch config
    'TIMESKETCH_HOST',
    'TIMESKETCH_USER',
    'TIMESKETCH_PASSWORD',]
# Environment variable to look for path data in
ENVCONFIGVAR = 'TURBINIA_CONFIG_PATH'

CONFIG = None

class TurbiniaConfigException(Exception):
  pass

def LoadConfig():
  """Finds Turbinia config file and loads it."""
  # TODO(aarontp): Find way to not require global var here.  Maybe a singleton
  # pattern on the config class.
  global CONFIG
  if CONFIG:
    return CONFIG

  if os.environ.has_key(ENVCONFIGVAR):
    CONFIGPATH.extend(os.environ[ENVCONFIGVAR].split(':'))

  config_file = None
  # Load first file found
  for _dir, _file in itertools.product(CONFIGPATH, CONFIGFILES):
    if os.path.exists(os.path.join(_dir, _file)):
      config_file = os.path.join(_dir, _file)
      break

  if config_file is None:
    raise TurbiniaConfigException('No config files found')

  log.info('Loading config from {0:s}'.format(config_file))
  _config = imp.load_source('config', config_file)
  _config.configSource = config_file
  ValidateAndSetConfig(_config)
  CONFIG = _config
  return _config

def ValidateAndSetConfig(_config):
  """Makes sure that the config has the vars loaded and set in the module."""
  # TODO(aarontp): Allow for non-mandatory config options
  for var in CONFIGVARS:
    if not hasattr(_config, var):
      raise TurbiniaConfigException(
          'No config attribute {0:s}:{1:s}'.format(_config.configSource, var))
    if getattr(_config, var) is None:
      raise TurbiniaConfigException(
          'Config attribute {0:s}:{1:s} is not set'.format(
              _config.configSource, var))

    # Set the attribute in the current module
    setattr(sys.modules[__name__], var, getattr(_config, var))
<commit_msg>Fix import typo (#141)<commit_after># -*- coding: utf-8 -*-
# Copyright 2016 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Basic Turbinia config."""

from __future__ import unicode_literals

import imp
import itertools
import logging
import os
import sys

log = logging.getLogger('turbinia')

# Look for config files with these names
CONFIGFILES = ['.turbiniarc', 'turbinia.conf', 'turbinia_config.py']
# Look in homedir first, then /etc/turbinia, and finally in the source
# config dir for config files
CONFIGPATH = [
    os.path.expanduser('~'),
    '/etc/turbinia',
    os.path.dirname(os.path.abspath(__file__))]
# Config vars that we expect to exist in the configuration
CONFIGVARS = [
    # Turbinia Config
    'TASK_MANAGER',
    'LOG_FILE',
    'OUTPUT_DIR',
    'SLEEP_TIME',
    'SINGLE_RUN',
    'MOUNT_DIR_PREFIX',
    'SHARED_FILESYSTEM',
    # GCE CONFIG
    'PROJECT',
    'ZONE',
    'INSTANCE',
    'DEVICE_NAME',
    'SCRATCH_PATH',
    'BUCKET_NAME',
    'PSQ_TOPIC',
    'PUBSUB_TOPIC',
    'GCS_OUTPUT_PATH',
    'STATE_MANAGER',
    # REDIS CONFIG
    'REDIS_HOST',
    'REDIS_PORT',
    # Timesketch config
    'TIMESKETCH_HOST',
    'TIMESKETCH_USER',
    'TIMESKETCH_PASSWORD',]
# Environment variable to look for path data in
ENVCONFIGVAR = 'TURBINIA_CONFIG_PATH'

CONFIG = None

class TurbiniaConfigException(Exception):
  pass

def LoadConfig():
  """Finds Turbinia config file and loads it."""
  # TODO(aarontp): Find way to not require global var here.  Maybe a singleton
  # pattern on the config class.
  global CONFIG
  if CONFIG:
    return CONFIG

  if os.environ.has_key(ENVCONFIGVAR):
    CONFIGPATH.extend(os.environ[ENVCONFIGVAR].split(':'))

  config_file = None
  # Load first file found
  for _dir, _file in itertools.product(CONFIGPATH, CONFIGFILES):
    if os.path.exists(os.path.join(_dir, _file)):
      config_file = os.path.join(_dir, _file)
      break

  if config_file is None:
    raise TurbiniaConfigException('No config files found')

  log.info('Loading config from {0:s}'.format(config_file))
  _config = imp.load_source('config', config_file)
  _config.configSource = config_file
  ValidateAndSetConfig(_config)
  CONFIG = _config
  return _config

def ValidateAndSetConfig(_config):
  """Makes sure that the config has the vars loaded and set in the module."""
  # TODO(aarontp): Allow for non-mandatory config options
  for var in CONFIGVARS:
    if not hasattr(_config, var):
      raise TurbiniaConfigException(
          'No config attribute {0:s}:{1:s}'.format(_config.configSource, var))
    if getattr(_config, var) is None:
      raise TurbiniaConfigException(
          'Config attribute {0:s}:{1:s} is not set'.format(
              _config.configSource, var))

    # Set the attribute in the current module
    setattr(sys.modules[__name__], var, getattr(_config, var))
<|endoftext|>


<commit_before>import json
import tempfile
import re
import datetime
import shutil
import zlib
from six.moves.urllib.request import urlopen
from exporters.readers.base_reader import BaseReader
from exporters.records.base_record import BaseRecord
from exporters.default_retries import retry_long, retry_short
from exporters.exceptions import ConfigurationError, InvalidDateRangeError
import logging

S3_URL_EXPIRES_IN = 1800  # half an hour should be enough

def get_bucket(bucket, aws_access_key_id, aws_secret_access_key, **kwargs):
    import boto

    if len(aws_access_key_id) > len(aws_secret_access_key):
        logging.warn("The AWS credential keys aren't in the usual size,"
                     " are you using the correct ones?")

    connection = boto.connect_s3(aws_access_key_id, aws_secret_access_key)
    try:
        return connection.get_bucket(bucket)
    except boto.exception.S3ResponseError:
        return connection.get_bucket(bucket, validate=False)

def format_prefixes(prefixes, start, end):
    import dateparser
    start_date = dateparser.parse(start or 'today')
    end_date = dateparser.parse(end or 'today')
    if start_date > end_date:
        raise InvalidDateRangeError

    dates = []
    while start_date <= end_date:
        dates.append(start_date)
        start_date += datetime.timedelta(days=1)

    return [date.strftime(p) for date in dates for p in prefixes]

def stream_decompress_multi(key):
    dec = zlib.decompressobj(16 + zlib.MAX_WBITS)
    while True:
        chunk = key.read(1024 * 8)
        if not chunk:
            break
        rv = dec.decompress(chunk)
        if rv:
            yield rv
        if dec.unused_data:
            unused = dec.unused_data
            dec = zlib.decompressobj(16 + zlib.MAX_WBITS)
            rv = dec.decompress(unused)
            if rv:
                yield rv

class S3BucketKeysFetcher(object):
    def __init__(self, reader_options, aws_access_key_id, aws_secret_access_key):
        self.source_bucket = get_bucket(
            reader_options.get('bucket'), aws_access_key_id, aws_secret_access_key)
        self.pattern = reader_options.get('pattern', None)

        prefix = reader_options.get('prefix', '')
        prefix_pointer = reader_options.get('prefix_pointer', '')
        prefix_format_using_date = reader_options.get('prefix_format_using_date')

        unformatted_prefixes = self._get_prefixes(prefix, prefix_pointer)
        try:
            start, end = self._get_prefix_formatting_dates(prefix_format_using_date)
        except ValueError:
            raise ConfigurationError('The option prefix_format_using_date '
                                     'should be either a date string or two '
                                     'date strings in a list/tuple')
        try:
            self.prefixes = format_prefixes(unformatted_prefixes, start, end)
        except InvalidDateRangeError:
            raise ConfigurationError('The end date should be greater or equal '
                                     'to the start date for the '
                                     'prefix_format_using_date option')

        self.logger = logging.getLogger('s3-reader')
        self.logger.setLevel(logging.INFO)

    def _get_prefixes(self, prefix, prefix_pointer):
        if prefix and prefix_pointer:
            raise ConfigurationError('prefix and prefix_pointer options '
                                     'cannot be used together')

        prefixes = [prefix] if isinstance(prefix, basestring) else prefix
        if prefix_pointer:
            prefixes = self._fetch_prefixes_from_pointer(prefix_pointer)
        return prefixes

    def _get_prefix_formatting_dates(self, prefix_dates):
        if not prefix_dates or isinstance(prefix_dates, basestring):
            prefix_dates = (prefix_dates, prefix_dates)
        return prefix_dates

    @retry_short
    def _download_pointer(self, prefix_pointer):
        return self.source_bucket.get_key(prefix_pointer).get_contents_as_string()

    def _fetch_prefixes_from_pointer(self, prefix_pointer):
        return filter(None, self._download_pointer(prefix_pointer).splitlines())

    def _get_keys_from_bucket(self):
        keys = []
        for prefix in self.prefixes:
            for key in self.source_bucket.list(prefix=prefix):
                if self.pattern:
                    if self._should_add_key(key):
                        keys.append(key.name)
                    else:
                        self.logger.info(
                            'Skipping S3 key {}. No match with pattern'.format(key.name))
                else:
                    keys.append(key.name)
        if self.pattern and not keys:
            self.logger.warn(
                'No S3 keys found that match provided pattern: {}'.format(self.pattern))
        return keys

    def _should_add_key(self, key):
        return bool(re.findall(self.pattern, key.name))

    def pending_keys(self):
        return self._get_keys_from_bucket()

class S3Reader(BaseReader):
    """
    Reads items from keys located in S3 buckets and compressed with gzip with a common path.

        - batch_size (int)
            Number of items to be returned in each batch

        - bucket (str)
            Name of the bucket to retrieve items from.

        - aws_access_key_id (str)
            Public access key to the s3 bucket.

        - aws_secret_access_key (str)
            Secret access key to the s3 bucket.

        - prefix (str)
            Prefix of s3 keys to be read.

        - prefix_pointer (str)
            Prefix pointing to the last version of dataset. This adds support for regular exports.
            For example:
                We have a weekly export set with CRON. If we wanted to point to a new data
                prefix every week, we should keep updating the export configuration. With a pointer,
                we can set the reader to read from that key, which contains one or several
                lines with valid prefixes to datasets, so only that pointer file should be updated.

        - pattern (str)
            S3 key name pattern (REGEX). All files that don't match this regex string will be
            discarded by the reader.

    """

    # List of options to set up the reader
    supported_options = {
        'batch_size': {'type': int, 'default': 10000},
        'bucket': {'type': basestring},
        'aws_access_key_id': {
            'type': basestring,
            'env_fallback': 'EXPORTERS_S3READER_AWS_KEY'
        },
        'aws_secret_access_key': {
            'type': basestring,
            'env_fallback': 'EXPORTERS_S3READER_AWS_SECRET'
        },
        'prefix': {'type': (basestring, list), 'default': ''},
        'prefix_pointer': {'type': basestring, 'default': None},
        'pattern': {'type': basestring, 'default': None},
        'prefix_format_using_date': {'type': (basestring, tuple, list), 'default': None}
    }

    def __init__(self, *args, **kwargs):
        super(S3Reader, self).__init__(*args, **kwargs)
        self.batch_size = self.read_option('batch_size')
        bucket_name = self.read_option('bucket')
        self.logger.info('Starting S3Reader for bucket: %s' % bucket_name)

        self.bucket = get_bucket(bucket_name,
                                 self.read_option('aws_access_key_id'),
                                 self.read_option('aws_secret_access_key'))

        self.keys_fetcher = S3BucketKeysFetcher(self.options,
                                                self.read_option('aws_access_key_id'),
                                                self.read_option('aws_secret_access_key'))
        self.keys = self.keys_fetcher.pending_keys()
        self.read_keys = []
        self.current_key = None
        self.last_block = 0
        self.logger.info('S3Reader has been initiated')
        self.tmp_folder = tempfile.mkdtemp()
        self.lines_reader = self.read_lines_from_keys()

    def get_read_streams(self):
        from exporters.bypasses.stream_bypass import Stream
        for key_name in self.keys:
            key = self.bucket.get_key(key_name)
            file_obj = urlopen(key.generate_url(S3_URL_EXPIRES_IN))
            yield Stream(file_obj, key_name, key.size)

    @retry_long
    def read_lines_from_keys(self):
        for current_key in self.keys:
            self.current_key = current_key
            self.last_position['current_key'] = current_key
            key = self.bucket.get_key(current_key)
            self.last_leftover = ''
            index_block = 0
            for uncompressed in stream_decompress_multi(key):
                if index_block >= self.last_block:
                    block_text = self.last_leftover + uncompressed
                    items = block_text.split('\n')
                    for i in items:
                        if i:
                            try:
                                object = json.loads(i)
                            except ValueError:
                                # Last uncomplete line
                                self.last_leftover = i
                                self.last_position['last_leftover'] = self.last_leftover
                            else:
                                item = BaseRecord(object)
                                yield item
                    self.last_block += 1
                index_block += 1

            self.read_keys.append(current_key)
            self.current_key = None
            self.last_position['keys'].remove(current_key)
            self.last_position['read_keys'] = self.read_keys
            self.last_position['current_key'] = self.current_key
            self.last_position['last_block'] = self.last_block
            self.last_block = 0

        self.finished = True

    def get_next_batch(self):
        """
        This method is called from the manager. It must return a list or a generator
        of BaseRecord objects.
        When it has nothing else to read, it must set class variable "finished" to True.
        """
        count = 0
        while count < self.batch_size:
            count += 1
            yield next(self.lines_reader)
        self.logger.debug('Done reading batch')

    def set_last_position(self, last_position):
        """
        Called from the manager, it is in charge of updating the last position of data commited
        by the writer, in order to have resume support
        """
        if last_position is None:
            self.last_position = {}
            self.last_position['keys'] = self.keys
            self.last_position['read_keys'] = self.read_keys
            self.last_position['current_key'] = None
            self.last_position['last_block'] = 0
        else:
            self.last_position = last_position
            self.keys = self.last_position['keys']
            self.read_keys = self.last_position['read_keys']
            if self.last_position['current_key']:
                self.current_key = self.last_position['current_key']
            else:
                self.current_key = self.keys[0]
                self.last_block = 0
            self.last_block = self.last_position['last_block']

    def close(self):
        shutil.rmtree(self.tmp_folder)
<commit_msg>Retry to read chunk<commit_after>import json
import tempfile
import re
import datetime
import shutil
import zlib
from six.moves.urllib.request import urlopen
from exporters.readers.base_reader import BaseReader
from exporters.records.base_record import BaseRecord
from exporters.default_retries import retry_long, retry_short
from exporters.exceptions import ConfigurationError, InvalidDateRangeError
import logging

S3_URL_EXPIRES_IN = 1800  # half an hour should be enough

def get_bucket(bucket, aws_access_key_id, aws_secret_access_key, **kwargs):
    import boto

    if len(aws_access_key_id) > len(aws_secret_access_key):
        logging.warn("The AWS credential keys aren't in the usual size,"
                     " are you using the correct ones?")

    connection = boto.connect_s3(aws_access_key_id, aws_secret_access_key)
    try:
        return connection.get_bucket(bucket)
    except boto.exception.S3ResponseError:
        return connection.get_bucket(bucket, validate=False)

def format_prefixes(prefixes, start, end):
    import dateparser
    start_date = dateparser.parse(start or 'today')
    end_date = dateparser.parse(end or 'today')
    if start_date > end_date:
        raise InvalidDateRangeError

    dates = []
    while start_date <= end_date:
        dates.append(start_date)
        start_date += datetime.timedelta(days=1)

    return [date.strftime(p) for date in dates for p in prefixes]

@retry_short
def read_chunk(key):
    return key.read(1024 * 8)

def stream_decompress_multi(key):
    dec = zlib.decompressobj(16 + zlib.MAX_WBITS)
    while True:
        chunk = read_chunk(key)
        if not chunk:
            break
        rv = dec.decompress(chunk)
        if rv:
            yield rv
        if dec.unused_data:
            unused = dec.unused_data
            dec = zlib.decompressobj(16 + zlib.MAX_WBITS)
            rv = dec.decompress(unused)
            if rv:
                yield rv

class S3BucketKeysFetcher(object):
    def __init__(self, reader_options, aws_access_key_id, aws_secret_access_key):
        self.source_bucket = get_bucket(
            reader_options.get('bucket'), aws_access_key_id, aws_secret_access_key)
        self.pattern = reader_options.get('pattern', None)

        prefix = reader_options.get('prefix', '')
        prefix_pointer = reader_options.get('prefix_pointer', '')
        prefix_format_using_date = reader_options.get('prefix_format_using_date')

        unformatted_prefixes = self._get_prefixes(prefix, prefix_pointer)
        try:
            start, end = self._get_prefix_formatting_dates(prefix_format_using_date)
        except ValueError:
            raise ConfigurationError('The option prefix_format_using_date '
                                     'should be either a date string or two '
                                     'date strings in a list/tuple')
        try:
            self.prefixes = format_prefixes(unformatted_prefixes, start, end)
        except InvalidDateRangeError:
            raise ConfigurationError('The end date should be greater or equal '
                                     'to the start date for the '
                                     'prefix_format_using_date option')

        self.logger = logging.getLogger('s3-reader')
        self.logger.setLevel(logging.INFO)

    def _get_prefixes(self, prefix, prefix_pointer):
        if prefix and prefix_pointer:
            raise ConfigurationError('prefix and prefix_pointer options '
                                     'cannot be used together')

        prefixes = [prefix] if isinstance(prefix, basestring) else prefix
        if prefix_pointer:
            prefixes = self._fetch_prefixes_from_pointer(prefix_pointer)
        return prefixes

    def _get_prefix_formatting_dates(self, prefix_dates):
        if not prefix_dates or isinstance(prefix_dates, basestring):
            prefix_dates = (prefix_dates, prefix_dates)
        return prefix_dates

    @retry_short
    def _download_pointer(self, prefix_pointer):
        return self.source_bucket.get_key(prefix_pointer).get_contents_as_string()

    def _fetch_prefixes_from_pointer(self, prefix_pointer):
        return filter(None, self._download_pointer(prefix_pointer).splitlines())

    def _get_keys_from_bucket(self):
        keys = []
        for prefix in self.prefixes:
            for key in self.source_bucket.list(prefix=prefix):
                if self.pattern:
                    if self._should_add_key(key):
                        keys.append(key.name)
                    else:
                        self.logger.info(
                            'Skipping S3 key {}. No match with pattern'.format(key.name))
                else:
                    keys.append(key.name)
        if self.pattern and not keys:
            self.logger.warn(
                'No S3 keys found that match provided pattern: {}'.format(self.pattern))
        return keys

    def _should_add_key(self, key):
        return bool(re.findall(self.pattern, key.name))

    def pending_keys(self):
        return self._get_keys_from_bucket()

class S3Reader(BaseReader):
    """
    Reads items from keys located in S3 buckets and compressed with gzip with a common path.

        - batch_size (int)
            Number of items to be returned in each batch

        - bucket (str)
            Name of the bucket to retrieve items from.

        - aws_access_key_id (str)
            Public access key to the s3 bucket.

        - aws_secret_access_key (str)
            Secret access key to the s3 bucket.

        - prefix (str)
            Prefix of s3 keys to be read.

        - prefix_pointer (str)
            Prefix pointing to the last version of dataset. This adds support for regular exports.
            For example:
                We have a weekly export set with CRON. If we wanted to point to a new data
                prefix every week, we should keep updating the export configuration. With a pointer,
                we can set the reader to read from that key, which contains one or several
                lines with valid prefixes to datasets, so only that pointer file should be updated.

        - pattern (str)
            S3 key name pattern (REGEX). All files that don't match this regex string will be
            discarded by the reader.

    """

    # List of options to set up the reader
    supported_options = {
        'batch_size': {'type': int, 'default': 10000},
        'bucket': {'type': basestring},
        'aws_access_key_id': {
            'type': basestring,
            'env_fallback': 'EXPORTERS_S3READER_AWS_KEY'
        },
        'aws_secret_access_key': {
            'type': basestring,
            'env_fallback': 'EXPORTERS_S3READER_AWS_SECRET'
        },
        'prefix': {'type': (basestring, list), 'default': ''},
        'prefix_pointer': {'type': basestring, 'default': None},
        'pattern': {'type': basestring, 'default': None},
        'prefix_format_using_date': {'type': (basestring, tuple, list), 'default': None}
    }

    def __init__(self, *args, **kwargs):
        super(S3Reader, self).__init__(*args, **kwargs)
        self.batch_size = self.read_option('batch_size')
        bucket_name = self.read_option('bucket')
        self.logger.info('Starting S3Reader for bucket: %s' % bucket_name)

        self.bucket = get_bucket(bucket_name,
                                 self.read_option('aws_access_key_id'),
                                 self.read_option('aws_secret_access_key'))

        self.keys_fetcher = S3BucketKeysFetcher(self.options,
                                                self.read_option('aws_access_key_id'),
                                                self.read_option('aws_secret_access_key'))
        self.keys = self.keys_fetcher.pending_keys()
        self.read_keys = []
        self.current_key = None
        self.last_block = 0
        self.logger.info('S3Reader has been initiated')
        self.tmp_folder = tempfile.mkdtemp()
        self.lines_reader = self.read_lines_from_keys()

    def get_read_streams(self):
        from exporters.bypasses.stream_bypass import Stream
        for key_name in self.keys:
            key = self.bucket.get_key(key_name)
            file_obj = urlopen(key.generate_url(S3_URL_EXPIRES_IN))
            yield Stream(file_obj, key_name, key.size)

    @retry_long
    def read_lines_from_keys(self):
        for current_key in self.keys:
            self.current_key = current_key
            self.last_position['current_key'] = current_key
            key = self.bucket.get_key(current_key)
            self.last_leftover = ''
            index_block = 0
            for uncompressed in stream_decompress_multi(key):
                if index_block >= self.last_block:
                    block_text = self.last_leftover + uncompressed
                    items = block_text.split('\n')
                    for i in items:
                        if i:
                            try:
                                object = json.loads(i)
                            except ValueError:
                                # Last uncomplete line
                                self.last_leftover = i
                                self.last_position['last_leftover'] = self.last_leftover
                            else:
                                item = BaseRecord(object)
                                yield item
                    self.last_block += 1
                index_block += 1

            self.read_keys.append(current_key)
            self.current_key = None
            self.last_position['keys'].remove(current_key)
            self.last_position['read_keys'] = self.read_keys
            self.last_position['current_key'] = self.current_key
            self.last_position['last_block'] = self.last_block
            self.last_block = 0

        self.finished = True

    def get_next_batch(self):
        """
        This method is called from the manager. It must return a list or a generator
        of BaseRecord objects.
        When it has nothing else to read, it must set class variable "finished" to True.
        """
        count = 0
        while count < self.batch_size:
            count += 1
            yield next(self.lines_reader)
        self.logger.debug('Done reading batch')

    def set_last_position(self, last_position):
        """
        Called from the manager, it is in charge of updating the last position of data commited
        by the writer, in order to have resume support
        """
        if last_position is None:
            self.last_position = {}
            self.last_position['keys'] = self.keys
            self.last_position['read_keys'] = self.read_keys
            self.last_position['current_key'] = None
            self.last_position['last_block'] = 0
        else:
            self.last_position = last_position
            self.keys = self.last_position['keys']
            self.read_keys = self.last_position['read_keys']
            if self.last_position['current_key']:
                self.current_key = self.last_position['current_key']
            else:
                self.current_key = self.keys[0]
                self.last_block = 0
            self.last_block = self.last_position['last_block']

    def close(self):
        shutil.rmtree(self.tmp_folder)
<|endoftext|>