<commit_before>import time
import pytest
import random
from threading import Timer
from Pyro4.errors import NamingError
from osbrain.core import locate_ns
from osbrain.core import run_agent
from osbrain.core import Agent
from osbrain.core import Proxy
from osbrain.core import NSProxy
from osbrain.core import NameServer
from osbrain.core import SocketAddress

@pytest.fixture(scope='function')
def nsaddr(request):
    while True:
        try:
            # Bind to random port
            host = '127.0.0.1'
            port = random.randrange(10000, 20000)
            addr = SocketAddress(host, port)
            ns = NameServer(addr)
            def terminate():
                print('addfinalizer...')
                ns.kill()
            request.addfinalizer(terminate)
            ns.start()
            return addr
        except NamingError:
            continue
        except PermissionError:
            continue

def test_nameserver(nsaddr):
    """
    A simple test that checks the correct creation of a name server.
    """
    nsproxy = NSProxy(nsaddr)
    agents = nsproxy.list()
    name = 'Pyro.NameServer'
    assert len(agents) == 1
    assert list(agents.keys())[0] == name
    assert agents[name] == 'PYRO:%s@%s' % (name, nsaddr)

def test_locate_ns():
    """
    Locate nameserver as fast as possible. The function `locate_ns` should
    have a timeout before raising an error.
    """
    while True:
        try:
            # Bind to random port
            host = '127.0.0.1'
            port = random.randrange(10000, 20000)
            addr = SocketAddress(host, port)
            ns = NameServer(addr)
            # Start name server later
            Timer(1, ns.start).start()
            # Locate name server now
            nsaddr = NSProxy(addr).addr()
        except PermissionError:
            continue
        break
    assert nsaddr.host == host
    assert nsaddr.port == port
    ns.kill()

def test_early_agent_proxy(nsaddr):
    """
    It must be possible to create a Proxy when the registration of the new
    agent is imminent, even if it has not occured yet. A timeout will occur
    in case the agent could not be located.
    """
    agent = Agent('a0', nsaddr)
    # Start agent later
    Timer(1, agent.start).start()
    # Locate agent now
    a0 = Proxy('a0', nsaddr)
    assert a0.test() == 'OK'

def test_agent_loopback(nsaddr):
    """
    An agent should always have a loopback inproc socket.
    """
    a0 = run_agent('a0', nsaddr)
    assert a0.get_addr('loopback') == 'inproc://loopback'

def test_ping(nsaddr):
    """
    Test simple loopback ping.
    """
    a0 = run_agent('a0', nsaddr)
    assert a0.ping() == 'PONG'

def test_registration(nsaddr):
    """
    Verify new agents get registered in the nameserver.
    """
    a0 = run_agent('a0', nsaddr)
    a1 = run_agent('a1', nsaddr)
    # List registered agents
    agent_list = NSProxy(nsaddr).list()
    assert 'a0' in agent_list
    assert 'a1' in agent_list

def test_socket_creation(nsaddr):
    """
    Test ZMQ socket creation.
    """
    a0 = run_agent('a0', nsaddr)
    a0.bind('REQ', 'alias0')
    a0.bind('PUB', 'alias1')
    a0.bind('PUSH', 'alias2')
    addresses = a0.get_attr('address')
    assert 'alias0' in addresses
    assert 'alias1' in addresses
    assert 'alias2' in addresses

# TODO: this function is used just within the scope of the next test.
#       Could we directly send the bytecode to the agent so that we can
#       declare it within a more constrained scope? (i.e. in the test code).
def rep_handler(agent, message):
    agent.send('reply', 'OK')

def test_reqrep(nsaddr):
    """
    Simple request-reply pattern between two agents.
    """
    a0 = run_agent('a0', nsaddr)
    a1 = run_agent('a1', nsaddr)
    addr = a0.bind('REP', 'reply', rep_handler)
    a1.connect(addr, 'request')
    response = a1.send_recv('request', 'Hello world')
    assert response == 'OK'
<commit_msg>Add push-pull pattern test<commit_after>import time
import pytest
import random
from threading import Timer
from Pyro4.errors import NamingError
from osbrain.core import locate_ns
from osbrain.core import run_agent
from osbrain.core import BaseAgent
from osbrain.core import Agent
from osbrain.core import Proxy
from osbrain.core import NSProxy
from osbrain.core import NameServer
from osbrain.core import SocketAddress

@pytest.fixture(scope='function')
def nsaddr(request):
    while True:
        try:
            # Bind to random port
            host = '127.0.0.1'
            port = random.randrange(10000, 20000)
            addr = SocketAddress(host, port)
            ns = NameServer(addr)
            def terminate():
                print('addfinalizer...')
                ns.kill()
            request.addfinalizer(terminate)
            ns.start()
            return addr
        except NamingError:
            continue
        except PermissionError:
            continue

def test_nameserver(nsaddr):
    """
    A simple test that checks the correct creation of a name server.
    """
    nsproxy = NSProxy(nsaddr)
    agents = nsproxy.list()
    name = 'Pyro.NameServer'
    assert len(agents) == 1
    assert list(agents.keys())[0] == name
    assert agents[name] == 'PYRO:%s@%s' % (name, nsaddr)

def test_locate_ns():
    """
    Locate nameserver as fast as possible. The function `locate_ns` should
    have a timeout before raising an error.
    """
    while True:
        try:
            # Bind to random port
            host = '127.0.0.1'
            port = random.randrange(10000, 20000)
            addr = SocketAddress(host, port)
            ns = NameServer(addr)
            # Start name server later
            Timer(1, ns.start).start()
            # Locate name server now
            nsaddr = NSProxy(addr).addr()
        except PermissionError:
            continue
        break
    assert nsaddr.host == host
    assert nsaddr.port == port
    ns.kill()

def test_early_agent_proxy(nsaddr):
    """
    It must be possible to create a Proxy when the registration of the new
    agent is imminent, even if it has not occured yet. A timeout will occur
    in case the agent could not be located.
    """
    agent = Agent('a0', nsaddr)
    # Start agent later
    Timer(1, agent.start).start()
    # Locate agent now
    a0 = Proxy('a0', nsaddr)
    assert a0.test() == 'OK'

def test_agent_loopback(nsaddr):
    """
    An agent should always have a loopback inproc socket.
    """
    a0 = run_agent('a0', nsaddr)
    assert a0.get_addr('loopback') == 'inproc://loopback'

def test_ping(nsaddr):
    """
    Test simple loopback ping.
    """
    a0 = run_agent('a0', nsaddr)
    assert a0.ping() == 'PONG'

def test_registration(nsaddr):
    """
    Verify new agents get registered in the nameserver.
    """
    a0 = run_agent('a0', nsaddr)
    a1 = run_agent('a1', nsaddr)
    # List registered agents
    agent_list = NSProxy(nsaddr).list()
    assert 'a0' in agent_list
    assert 'a1' in agent_list

def test_socket_creation(nsaddr):
    """
    Test ZMQ socket creation.
    """
    a0 = run_agent('a0', nsaddr)
    a0.bind('REQ', 'alias0')
    a0.bind('PUB', 'alias1')
    a0.bind('PUSH', 'alias2')
    addresses = a0.get_attr('address')
    assert 'alias0' in addresses
    assert 'alias1' in addresses
    assert 'alias2' in addresses

# TODO: this functions are used just within the scope of the next test.
#       Could we directly send the bytecode to the agent so that we can
#       declare them within a more constrained scope? (i.e. in the test code).
def rep_handler(agent, message):
    agent.send('reply', 'OK')
def redirect(agent, message):
    agent.send('push', '%s (redirected)' % message)
def set_received(agent, message):
    agent.received = message

def test_reqrep(nsaddr):
    """
    Simple request-reply pattern between two agents.
    """
    a0 = run_agent('a0', nsaddr)
    a1 = run_agent('a1', nsaddr)
    addr = a0.bind('REP', 'reply', rep_handler)
    a1.connect(addr, 'request')
    response = a1.send_recv('request', 'Hello world')
    assert response == 'OK'

def test_pushpull(nsaddr):
    """
    Simple push-pull pattern between two agents.
    """
    time.sleep(2)
    a0 = run_agent('a0', nsaddr)
    a1 = run_agent('a1', nsaddr)
    addr = a1.bind('PULL', 'pull', redirect)
    a0.connect(addr, 'push')
    # Create a BaseAgent as end-point
    a2 = BaseAgent('a2')
    a2.received = ''
    addr = a2.bind('PULL', 'pull', set_received)
    a1.connect(addr, 'push')
    # Send message (will be passed from a0 to a1 and then to a2)
    message = 'Hello world'
    a0.send('push', message)
    while not a2.received:
        a2.iterate()
    assert a2.received == '%s (redirected)' % message
<|endoftext|>


<commit_before>#!/usr/bin/python3

'''
gdb python script

notes:
    - you should avoid running manually the same target program while this is running, get_process_pid will not like it
'''

import subprocess
import optparse
import threading
import psutil
import time
import os
import sys

####GLOBAL###
exePath="/usr/lib/libreoffice/program/soffice.bin"
exeArgs='--impress --norestore'
#testcasesPath="/mnt/shared/ppt"
crashesPath=""
logsPath=""
fuzzerPath="./radamsa-bin"
fuzzIter=200
fuzzDst="fuzzed"
restoreFlag = False
debugFlag = True
processTimeout=20 #seconds
listfile='filelist.txt'

#############

"""
for every file
fuzz it
iterate fuzzed
check for crashes
close program if nothing happens
"""

def die(msg):
    sys.stderr.write("[ERROR] " + msg + "\n")
    sys.exit(-1)

def debug_msg(msg):
    global debugFlag
    if debugFlag:
        sys.stdout.write('[FUZZER] ' + msg + '\n')
    return

#receive full path of testcase, and dst dir
#write target files
def fuzz_testcase(testcase, fuzzDst):
    global fuzzerPath
    global fuzzIter
    if not os.path.exists(fuzzDst):
        debug_msg("creating dir %s\n" % (fuzzDst))
        os.mkdir(fuzzDst)
    if debugFlag:
        fuzzCmd = "%s -v -n %d -o %s/fuzzed-%%n.ppt %s" % (fuzzerPath, fuzzIter, fuzzDst, testcase)
        debug_msg(fuzzCmd)
    else:
        fuzzCmd = "%s -n %d -o %s/fuzzed-%%n.ppt %s" % (fuzzerPath, fuzzIter, fuzzDst, testcase)
    ret = os.system(fuzzCmd)
    if (ret != 0):
        die("fuzzer failed to run")
    return

def empty_fuzzdir(fuzzDst):
    os.system("rm -rf %s/*" % fuzzDst)

def parse_args():
    parser = optparse.OptionParser("%prog [some opts] [-L filelist]|[-D fuzzdir]")
    parser.add_option("-v", "--debug", help="get debug output", action="store_true", dest="debug", default=True)
    parser.add_option("-s", "--export", help="save filelist.txt ", action="store_true", dest="saveList", default=False)
    parser.add_option("-t", "--nofuzz", help="start fuzzer, dont generate testcases, for TESTING", action="store_true", dest="nofuzz", default=False)
    parser.add_option("-T", "--test", help="don't start fuzzer, for TESTING", action="store_true", dest="testrun", default=False)
    parser.add_option("-S", "--skipto", help="skip to #n iteration", dest="skipto", default=None)
    parser.add_option("-L", "--list", help="read filelist from file", dest="filelist", default=None)
    parser.add_option("-D", "--fuzzdir", help="create filelist from dir", dest="fuzzdir", default=None)

    return parser.parse_args()

def main():
    global exePath
    global exeArgs
    global testcasesPath
    global fuzzDst
    global debugFlag
    global listfile

    opts, args = parse_args()

    debugFlag = opts.debug

    if opts.filelist and opts.fuzzdir:
        die("options -L and -D are mutually exclusive")
    if not opts.filelist and not opts.fuzzdir:
        die("dickhead")
    if opts.fuzzdir:
        testcasesPath=opts.fuzzdir
        filelist=os.listdir(testcasesPath)
        filelist=[testcasePath+'/'+ tc for tc in filelist]
    if opts.filelist:
        if opts.saveList:
            debug_msg("already loading from filelist, unsetting writing")
            opts.saveList = False
        try:
            lf=open(opts.filelist)
            filelist=lf.read().split('\n')
        except:
            die('Impossible to load file list, check file and syntax')

    if not filelist:
        die("no input files were specified")

    if opts.saveList:
        debug_msg("saving list of files to %s" % listfile)
        lf = open(listfile, 'w')
        lf.write("\n".join([str(i) for i in filelist]))
        lf.close()

    if opts.testrun:
        debug_msg("testrun only, exiting")
        sys.exit(0)

    #XXX: why is this try here?
    try:
        start=0
        if opts.skipto:
            start=int(opts.skipto)
        for i in range(start,len(filelist)):
            f = filelist[i]
            fpfile="%s/%s" % (testcasesPath, f)
            debug_msg('fuzzing testcase #%d : %s' % (i,fpfile))
            if not opts.nofuzz:
                fuzz_testcase(fpfile, fuzzDst)
            for file in os.listdir(fuzzDst):
                fuzzedcase=fuzzDst + "/" + file
                debug_msg("run target with file %s" % fuzzedcase)
                gdb_proc = subprocess.Popen("./launcher.py --batch --args %s %s %s" % (exePath, exeArgs, fuzzedcase), shell="/usr/bin/python")
                mon_proc = subprocess.Popen("./process_monitor.py %s %s" % (exePath, fuzzedcase), shell="/usr/bin/python")
                if not opts.nofuzz:
                    debug_msg("nofuzz set, will not destroy testcases")
                else:
                    empty_fuzzdir(fuzzDst)
                gdb_proc.wait()
                mon_proc.wait()
                debug_msg('Terminated fuzzing %s' % fuzzedcase)
            if opts.nofuzz:
                debug_msg("nofuzz set, will only do first iteration")
                #gdb_proc.kill()
                #mon_proc.kill()
                sys.exit(0)
    except KeyboardInterrupt:
        debug_msg("Ctrl-c detected, exiting")
        try:
            gdb_proc.kill()
            mon_proc.kill()
        except:
            pass

        sys.exit(0)

main()
<commit_msg>fixes<commit_after>#!/usr/bin/python3

'''
gdb python script

notes:
    - you should avoid running manually the same target program while this is running, get_process_pid will not like it
'''

import subprocess
import optparse
import threading
import psutil
import time
import os
import sys

####GLOBAL###
exePath="/usr/lib/libreoffice/program/soffice.bin"
exeArgs='--impress --norestore'
#testcasesPath="/mnt/shared/ppt"
crashesPath=""
logsPath=""
fuzzerPath="./radamsa-bin"
fuzzIter=200
fuzzDst="fuzzed"
restoreFlag = False
debugFlag = True
processTimeout=20 #seconds
listfile='filelist.txt'

#############

"""
for every file
fuzz it
iterate fuzzed
check for crashes
close program if nothing happens
"""

def die(msg):
    sys.stderr.write("[ERROR] " + msg + "\n")
    sys.exit(-1)

def debug_msg(msg):
    global debugFlag
    if debugFlag:
        sys.stdout.write('[FUZZER] ' + msg + '\n')
    return

#receive full path of testcase, and dst dir
#write target files
def fuzz_testcase(testcase, fuzzDst):
    global fuzzerPath
    global fuzzIter
    if not os.path.exists(fuzzDst):
        debug_msg("creating dir %s\n" % (fuzzDst))
        os.mkdir(fuzzDst)
    if debugFlag:
        fuzzCmd = "%s -v -n %d -o %s/fuzzed-%%n.ppt %s" % (fuzzerPath, fuzzIter, fuzzDst, testcase)
        debug_msg(fuzzCmd)
    else:
        fuzzCmd = "%s -n %d -o %s/fuzzed-%%n.ppt %s" % (fuzzerPath, fuzzIter, fuzzDst, testcase)
    ret = os.system(fuzzCmd)
    if (ret != 0):
        die("fuzzer failed to run")
    return

def empty_fuzzdir(fuzzDst):
    os.system("rm -rf %s/*" % fuzzDst)

def parse_args():
    parser = optparse.OptionParser("%prog [some opts] [-L filelist]|[-D fuzzdir]")
    parser.add_option("-v", "--debug", help="get debug output", action="store_true", dest="debug", default=True)
    parser.add_option("-s", "--export", help="save filelist.txt ", action="store_true", dest="saveList", default=False)
    parser.add_option("-t", "--nofuzz", help="start fuzzer, dont generate testcases, for TESTING", action="store_true", dest="nofuzz", default=False)
    parser.add_option("-T", "--test", help="don't start fuzzer, for TESTING", action="store_true", dest="testrun", default=False)
    parser.add_option("-S", "--skipto", help="skip to #n iteration", dest="skipto", default=None)
    parser.add_option("-L", "--list", help="read filelist from file", dest="filelist", default=None)
    parser.add_option("-D", "--fuzzdir", help="create filelist from dir", dest="fuzzdir", default=None)

    return parser.parse_args()

def main():
    global exePath
    global exeArgs
    global testcasesPath
    global fuzzDst
    global debugFlag
    global listfile

    opts, args = parse_args()

    debugFlag = opts.debug

    if opts.filelist and opts.fuzzdir:
        die("options -L and -D are mutually exclusive")
    if not opts.filelist and not opts.fuzzdir:
        die("dickhead")
    if opts.fuzzdir:
        testcasesPath=opts.fuzzdir
        filelist=os.listdir(testcasesPath)
        filelist=[testcasesPath+'/'+ tc for tc in filelist]
    if opts.filelist:
        if opts.saveList:
            debug_msg("already loading from filelist, unsetting writing")
            opts.saveList = False
        try:
            lf=open(opts.filelist)
            filelist=lf.read().split('\n')
        except:
            die('Impossible to load file list, check file and syntax')

    if not filelist:
        die("no input files were specified")

    if opts.saveList:
        debug_msg("saving list of files to %s" % listfile)
        lf = open(listfile, 'w')
        lf.write("\n".join([str(i) for i in filelist]))
        lf.close()

    if opts.testrun:
        debug_msg("testrun only, exiting")
        sys.exit(0)

    #XXX: why is this try here?
    try:
        start=0
        if opts.skipto:
            start=int(opts.skipto)
        for i in range(start,len(filelist)):
            f = filelist[i]
            debug_msg('fuzzing testcase #%d : %s' % (i,fpfile))
            if not opts.nofuzz:
                fuzz_testcase(f, fuzzDst)
            for file in os.listdir(fuzzDst):
                fuzzedcase=fuzzDst + "/" + file
                debug_msg("run target with file %s" % fuzzedcase)
                gdb_proc = subprocess.Popen("./launcher.py --batch --args %s %s %s" % (exePath, exeArgs, fuzzedcase), shell="/usr/bin/python")
                mon_proc = subprocess.Popen("./process_monitor.py %s %s" % (exePath, fuzzedcase), shell="/usr/bin/python")
                if not opts.nofuzz:
                    debug_msg("nofuzz set, will not destroy testcases")
                else:
                    empty_fuzzdir(fuzzDst)
                gdb_proc.wait()
                mon_proc.wait()
                debug_msg('Terminated fuzzing %s' % fuzzedcase)
            if opts.nofuzz:
                debug_msg("nofuzz set, will only do first iteration")
                #gdb_proc.kill()
                #mon_proc.kill()
                sys.exit(0)
    except KeyboardInterrupt:
        debug_msg("Ctrl-c detected, exiting")
        try:
            gdb_proc.kill()
            mon_proc.kill()
        except:
            pass

        sys.exit(0)

main()
<|endoftext|>


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

from blinkenlights import dimmer, setup, cleanup

pin = 18

setup(pin)

up = range(10)
down = range(9)
down.reverse()

bpm = 70

period = 60.0 / bpm # seconds

time_per_level = period / len(up + down)

for i in range(10):
    for j in (up + down):
        brightness = (j+1) / 10.0
        dimmer(brightness, time_per_level, pin)

cleanup()
<commit_msg>Add a 0.0 duty cycle: brief moment of off time.<commit_after>#!/usr/bin/env python

from blinkenlights import dimmer, setup, cleanup

pin = 18
bpm = 70

setup(pin)

up = range(10)
down = range(9)
down.reverse()
spectrum = up + down + [-1]

period = 60.0 / bpm # seconds
time_per_level = period / len(spectrum)

for i in range(10):
    for j in (spectrum):
        brightness = (j+1) / 10.0
        dimmer(brightness, time_per_level, pin)

cleanup()
<|endoftext|>


<commit_before>"""Manages logger for dvc repo."""

from __future__ import unicode_literals

from dvc.exceptions import DvcException
from dvc.utils.compat import str
from dvc.progress import progress_aware

import re
import sys
import logging
import traceback

from contextlib import contextmanager

import colorama

@progress_aware
def info(message):
    """Prints an info message."""
    logger.info(message)

def debug(message):
    """Prints a debug message."""
    prefix = colorize("Debug", color="blue")

    out = "{prefix}: {message}".format(prefix=prefix, message=message)

    logger.debug(out)

@progress_aware
def warning(message, parse_exception=False):
    """Prints a warning message."""
    prefix = colorize("Warning", color="yellow")

    exception, stack_trace = None, ""
    if parse_exception:
        exception, stack_trace = _parse_exc()

    out = "{prefix}: {description}".format(
        prefix=prefix, description=_description(message, exception)
    )

    if stack_trace:
        out += "\n{stack_trace}".format(stack_trace=stack_trace)

    logger.warning(out)

@progress_aware
def error(message=None):
    """Prints an error message."""
    prefix = colorize("Error", color="red")

    exception, stack_trace = _parse_exc()

    out = (
        "{prefix}: {description}"
        "\n"
        "{stack_trace}"
        "\n"
        "{footer}".format(
            prefix=prefix,
            description=_description(message, exception),
            stack_trace=stack_trace,
            footer=_footer(),
        )
    )

    logger.error(out)

def box(message, border_color=None):
    """Prints a message in a box.

    Args:
        message (unicode): message to print.
        border_color (unicode): name of a color to outline the box with.
    """
    lines = message.split("\n")
    max_width = max(_visual_width(line) for line in lines)

    padding_horizontal = 5
    padding_vertical = 1

    box_size_horizontal = max_width + (padding_horizontal * 2)

    chars = {"corner": "+", "horizontal": "-", "vertical": "|", "empty": " "}

    margin = "{corner}{line}{corner}\n".format(
        corner=chars["corner"], line=chars["horizontal"] * box_size_horizontal
    )

    padding_lines = [
        "{border}{space}{border}\n".format(
            border=colorize(chars["vertical"], color=border_color),
            space=chars["empty"] * box_size_horizontal,
        )
        * padding_vertical
    ]

    content_lines = [
        "{border}{space}{content}{space}{border}\n".format(
            border=colorize(chars["vertical"], color=border_color),
            space=chars["empty"] * padding_horizontal,
            content=_visual_center(line, max_width),
        )
        for line in lines
    ]

    box_str = "{margin}{padding}{content}{padding}{margin}".format(
        margin=colorize(margin, color=border_color),
        padding="".join(padding_lines),
        content="".join(content_lines),
    )

    logger.info(box_str)

def level():
    """Returns current log level."""
    return logger.getEffectiveLevel()

def set_level(level_name):
    """Sets log level.

    Args:
        level_name (str): log level name. E.g. info, debug, warning, error,
            critical.
    """
    if not level_name:
        return

    levels = {
        "info": logging.INFO,
        "debug": logging.DEBUG,
        "warning": logging.WARNING,
        "error": logging.ERROR,
        "critical": logging.CRITICAL,
    }

    logger.setLevel(levels.get(level_name))

def be_quiet():
    """Disables all messages except critical ones."""
    logger.setLevel(logging.CRITICAL)

def be_verbose():
    """Enables all messages."""
    logger.setLevel(logging.DEBUG)

@contextmanager
def verbose():
    """Enables verbose mode for the context."""
    previous_level = level()
    be_verbose()
    yield
    logger.setLevel(previous_level)

@contextmanager
def quiet():
    """Enables quiet mode for the context."""
    previous_level = level()
    be_quiet()
    yield
    logger.setLevel(previous_level)

def is_quiet():
    """Returns whether or not all messages except critical ones are
    disabled.
    """
    return level() == logging.CRITICAL

def is_verbose():
    """Returns whether or not all messages are enabled."""
    return level() == logging.DEBUG

def colorize(message, color=None):
    """Returns a message in a specified color."""
    if not color:
        return message

    colors = {
        "green": colorama.Fore.GREEN,
        "yellow": colorama.Fore.YELLOW,
        "blue": colorama.Fore.BLUE,
        "red": colorama.Fore.RED,
    }

    return "{color}{message}{nc}".format(
        color=colors.get(color, ""), message=message, nc=colorama.Fore.RESET
    )

def _init_colorama():
    colorama.init()

def set_default_level():
    """Sets default log level."""
    logger.setLevel(logging.INFO)

def _add_handlers():
    formatter = "%(message)s"

    class _LogLevelFilter(logging.Filter):
        # pylint: disable=too-few-public-methods
        def filter(self, record):
            return record.levelno <= logging.WARNING

    sh_out = logging.StreamHandler(sys.stdout)
    sh_out.setFormatter(logging.Formatter(formatter))
    sh_out.setLevel(logging.DEBUG)
    sh_out.addFilter(_LogLevelFilter())

    sh_err = logging.StreamHandler(sys.stderr)
    sh_err.setFormatter(logging.Formatter(formatter))
    sh_err.setLevel(logging.ERROR)

    logger.addHandler(sh_out)
    logger.addHandler(sh_err)

def _walk_exc(exc):
    exc_list = [str(exc)]
    tb_list = [traceback.format_exc()]

    # NOTE: parsing chained exceptions. See dvc/exceptions.py for more info.
    while hasattr(exc, "cause") and exc.cause is not None:
        exc_list.append(str(exc.cause))
        if hasattr(exc, "cause_tb") and exc.cause_tb is not None:
            tb_list.insert(0, str(exc.cause_tb))
        exc = exc.cause

    return exc_list, tb_list

def _parse_exc():
    exc = sys.exc_info()[1]
    if not exc:
        return (None, "")

    exc_list, tb_list = _walk_exc(exc)

    exception = ": ".join(exc_list)

    if is_verbose():
        stack_trace = "{line}\n{stack_trace}{line}\n".format(
            line=colorize("-" * 60, color="red"),
            stack_trace="\n".join(tb_list),
        )
    else:
        stack_trace = ""

    return (exception, stack_trace)

def _description(message, exception):
    if exception and message:
        description = "{message} - {exception}"
    elif exception:
        description = "{exception}"
    elif message:
        description = "{message}"
    else:
        raise DvcException(
            "Unexpected error - either exception or message must be provided"
        )

    return description.format(message=message, exception=exception)

def _footer():
    return "{phrase} Hit us up at {url}, we are always happy to help!".format(
        phrase=colorize("Having any troubles?", "yellow"),
        url=colorize("https://dvc.org/support", "blue"),
    )

def _visual_width(line):
    """Get the the number of columns required to display a string"""

    return len(re.sub(colorama.ansitowin32.AnsiToWin32.ANSI_CSI_RE, "", line))

def _visual_center(line, width):
    """Center align string according to it's visual width"""

    spaces = max(width - _visual_width(line), 0)
    left_padding = int(spaces / 2)
    right_padding = spaces - left_padding

    return (left_padding * " ") + line + (right_padding * " ")

logger = logging.getLogger("dvc")  # pylint: disable=invalid-name

set_default_level()
_add_handlers()
_init_colorama()
<commit_msg>cli: Use correct color codes on Windows<commit_after>"""Manages logger for dvc repo."""

from __future__ import unicode_literals

from dvc.exceptions import DvcException
from dvc.utils.compat import str
from dvc.progress import progress_aware

import re
import sys
import logging
import traceback

from contextlib import contextmanager

import colorama

@progress_aware
def info(message):
    """Prints an info message."""
    logger.info(message)

def debug(message):
    """Prints a debug message."""
    prefix = colorize("Debug", color="blue")

    out = "{prefix}: {message}".format(prefix=prefix, message=message)

    logger.debug(out)

@progress_aware
def warning(message, parse_exception=False):
    """Prints a warning message."""
    prefix = colorize("Warning", color="yellow")

    exception, stack_trace = None, ""
    if parse_exception:
        exception, stack_trace = _parse_exc()

    out = "{prefix}: {description}".format(
        prefix=prefix, description=_description(message, exception)
    )

    if stack_trace:
        out += "\n{stack_trace}".format(stack_trace=stack_trace)

    logger.warning(out)

@progress_aware
def error(message=None):
    """Prints an error message."""
    prefix = colorize("Error", color="red")

    exception, stack_trace = _parse_exc()

    out = (
        "{prefix}: {description}"
        "\n"
        "{stack_trace}"
        "\n"
        "{footer}".format(
            prefix=prefix,
            description=_description(message, exception),
            stack_trace=stack_trace,
            footer=_footer(),
        )
    )

    logger.error(out)

def box(message, border_color=None):
    """Prints a message in a box.

    Args:
        message (unicode): message to print.
        border_color (unicode): name of a color to outline the box with.
    """
    lines = message.split("\n")
    max_width = max(_visual_width(line) for line in lines)

    padding_horizontal = 5
    padding_vertical = 1

    box_size_horizontal = max_width + (padding_horizontal * 2)

    chars = {"corner": "+", "horizontal": "-", "vertical": "|", "empty": " "}

    margin = "{corner}{line}{corner}\n".format(
        corner=chars["corner"], line=chars["horizontal"] * box_size_horizontal
    )

    padding_lines = [
        "{border}{space}{border}\n".format(
            border=colorize(chars["vertical"], color=border_color),
            space=chars["empty"] * box_size_horizontal,
        )
        * padding_vertical
    ]

    content_lines = [
        "{border}{space}{content}{space}{border}\n".format(
            border=colorize(chars["vertical"], color=border_color),
            space=chars["empty"] * padding_horizontal,
            content=_visual_center(line, max_width),
        )
        for line in lines
    ]

    box_str = "{margin}{padding}{content}{padding}{margin}".format(
        margin=colorize(margin, color=border_color),
        padding="".join(padding_lines),
        content="".join(content_lines),
    )

    logger.info(box_str)

def level():
    """Returns current log level."""
    return logger.getEffectiveLevel()

def set_level(level_name):
    """Sets log level.

    Args:
        level_name (str): log level name. E.g. info, debug, warning, error,
            critical.
    """
    if not level_name:
        return

    levels = {
        "info": logging.INFO,
        "debug": logging.DEBUG,
        "warning": logging.WARNING,
        "error": logging.ERROR,
        "critical": logging.CRITICAL,
    }

    logger.setLevel(levels.get(level_name))

def be_quiet():
    """Disables all messages except critical ones."""
    logger.setLevel(logging.CRITICAL)

def be_verbose():
    """Enables all messages."""
    logger.setLevel(logging.DEBUG)

@contextmanager
def verbose():
    """Enables verbose mode for the context."""
    previous_level = level()
    be_verbose()
    yield
    logger.setLevel(previous_level)

@contextmanager
def quiet():
    """Enables quiet mode for the context."""
    previous_level = level()
    be_quiet()
    yield
    logger.setLevel(previous_level)

def is_quiet():
    """Returns whether or not all messages except critical ones are
    disabled.
    """
    return level() == logging.CRITICAL

def is_verbose():
    """Returns whether or not all messages are enabled."""
    return level() == logging.DEBUG

def colorize(message, color=None):
    """Returns a message in a specified color."""
    if not color:
        return message

    colors = {
        "green": colorama.Fore.GREEN,
        "yellow": colorama.Fore.YELLOW,
        "blue": colorama.Fore.BLUE,
        "red": colorama.Fore.RED,
    }

    return "{color}{message}{nc}".format(
        color=colors.get(color, ""), message=message, nc=colorama.Fore.RESET
    )

def set_default_level():
    """Sets default log level."""
    logger.setLevel(logging.INFO)

def _add_handlers():
    # NOTE: We need to initialize colorama before setting the stream handlers
    #       so it can wrap stdout/stderr and convert color codes to Windows.
    colorama.init()

    formatter = "%(message)s"

    class _LogLevelFilter(logging.Filter):
        # pylint: disable=too-few-public-methods
        def filter(self, record):
            return record.levelno <= logging.WARNING

    sh_out = logging.StreamHandler(sys.stdout)
    sh_out.setFormatter(logging.Formatter(formatter))
    sh_out.setLevel(logging.DEBUG)
    sh_out.addFilter(_LogLevelFilter())

    sh_err = logging.StreamHandler(sys.stderr)
    sh_err.setFormatter(logging.Formatter(formatter))
    sh_err.setLevel(logging.ERROR)

    logger.addHandler(sh_out)
    logger.addHandler(sh_err)

def _walk_exc(exc):
    exc_list = [str(exc)]
    tb_list = [traceback.format_exc()]

    # NOTE: parsing chained exceptions. See dvc/exceptions.py for more info.
    while hasattr(exc, "cause") and exc.cause is not None:
        exc_list.append(str(exc.cause))
        if hasattr(exc, "cause_tb") and exc.cause_tb is not None:
            tb_list.insert(0, str(exc.cause_tb))
        exc = exc.cause

    return exc_list, tb_list

def _parse_exc():
    exc = sys.exc_info()[1]
    if not exc:
        return (None, "")

    exc_list, tb_list = _walk_exc(exc)

    exception = ": ".join(exc_list)

    if is_verbose():
        stack_trace = "{line}\n{stack_trace}{line}\n".format(
            line=colorize("-" * 60, color="red"),
            stack_trace="\n".join(tb_list),
        )
    else:
        stack_trace = ""

    return (exception, stack_trace)

def _description(message, exception):
    if exception and message:
        description = "{message} - {exception}"
    elif exception:
        description = "{exception}"
    elif message:
        description = "{message}"
    else:
        raise DvcException(
            "Unexpected error - either exception or message must be provided"
        )

    return description.format(message=message, exception=exception)

def _footer():
    return "{phrase} Hit us up at {url}, we are always happy to help!".format(
        phrase=colorize("Having any troubles?", "yellow"),
        url=colorize("https://dvc.org/support", "blue"),
    )

def _visual_width(line):
    """Get the the number of columns required to display a string"""

    return len(re.sub(colorama.ansitowin32.AnsiToWin32.ANSI_CSI_RE, "", line))

def _visual_center(line, width):
    """Center align string according to it's visual width"""

    spaces = max(width - _visual_width(line), 0)
    left_padding = int(spaces / 2)
    right_padding = spaces - left_padding

    return (left_padding * " ") + line + (right_padding * " ")

logger = logging.getLogger("dvc")  # pylint: disable=invalid-name

set_default_level()
_add_handlers()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import boto.dynamodb2.layer1, json, sys, time, shutil, os, argparse, logging, datetime, threading
from boto.dynamodb2.layer1 import DynamoDBConnection

JSON_INDENT = 2
AWS_SLEEP_INTERVAL = 10 #seconds
LOCAL_SLEEP_INTERVAL = 1 #seconds
MAX_BATCH_WRITE = 25 #DynamoDB limit
SCHEMA_FILE = "schema.json"
DATA_DIR = "data"
MAX_RETRY = 3
LOCAL_REGION = "local"
LOG_LEVEL = "INFO"
DUMP_PATH = "dump"
RESTORE_WRITE_CAPACITY = 100
THREAD_START_DELAY = 1 #seconds

def get_table_name_matches(conn, table_name_wildcard):
  all_tables = []
  last_evaluated_table_name = None

  while True:
    table_list = conn.list_tables(exclusive_start_table_name=last_evaluated_table_name)
    all_tables.extend(table_list["TableNames"])

    try:
      last_evaluated_table_name = table_list["LastEvaluatedTableName"]
    except KeyError, e:
      break

  matching_tables = []
  for table_name in all_tables:
    if table_name.startswith(table_name_wildcard.split("*", 1)[0]):
      matching_tables.append(table_name)

  return matching_tables

def get_restore_table_matches(table_name_wildcard):
  matching_tables = []
  dir_list = os.listdir("./" + DUMP_PATH)
  for dir_name in dir_list:
    if dir_name.startswith(table_name_wildcard.split("*", 1)[0]):
      matching_tables.append(dir_name)

  return matching_tables

def change_prefix(source_table_name, source_wildcard, destination_wildcard):
  source_prefix = source_wildcard.split("*", 1)[0]
  destination_prefix = destination_wildcard.split("*", 1)[0]
  if source_table_name.split("-", 1)[0] == source_prefix:
    return destination_prefix + "-" + source_table_name.split("-", 1)[1]

def delete_table(conn, sleep_interval, table_name):
  # delete table if exists
  table_exist = True
  try:
    conn.delete_table(table_name)
  except boto.exception.JSONResponseError, e:
    if e.body["__type"] == "com.amazonaws.dynamodb.v20120810#ResourceNotFoundException":
      table_exist = False
      logging.info("Table does not exist in destination, deleting not necessary..")
      pass
    else:
      logging.exception(e)
      sys.exit(1)

  # if table exists, wait till deleted
  if table_exist:
    try:
      while True:
        logging.info("Waiting for " + table_name + " table to be deleted.. [" + conn.describe_table(table_name)["Table"]["TableStatus"] +"]")
        time.sleep(sleep_interval)
    except boto.exception.JSONResponseError, e:
      if e.body["__type"] == "com.amazonaws.dynamodb.v20120810#ResourceNotFoundException":
        logging.info(table_name + " table deleted.")
        pass
      else:
        logging.exception(e)
        sys.exit(1)

def mkdir_p(path):
  try:
    os.makedirs(path)
  except OSError as exc:
    if exc.errno == errno.EEXIST and os.path.isdir(path):
      pass
    else: raise

def batch_write(conn, sleep_interval, table_name, put_requests):
  request_items = {table_name: put_requests}
  i = 1
  while True:
    response = conn.batch_write_item(request_items)
    unprocessed_items = response["UnprocessedItems"]

    if len(unprocessed_items) == 0:
      break

    if len(unprocessed_items) > 0 and i <= MAX_RETRY:
      logging.info(str(len(unprocessed_items)) + " unprocessed items, retrying.. [" + str(i) + "]")
      request_items = unprocessed_items
      i += 1
      time.sleep(sleep_interval)
    else:
      logging.info("Max retries reached, failed to processed batch write: " + json.dumps(unprocessed_items, indent=JSON_INDENT))
      logging.info("Ignoring and continuing..")
      break

def wait_for_active_table(conn, table_name, verb):
  while True:
    if conn.describe_table(table_name)["Table"]["TableStatus"] != "ACTIVE":
      logging.info("Waiting for " + table_name + " table to be " + verb + ".. [" + conn.describe_table(table_name)["Table"]["TableStatus"] +"]")
      time.sleep(sleep_interval)
    else:
      logging.info(table_name + " " + verb + ".")
      break

def update_provisioned_throughput(conn, table_name, read_capacity, write_capacity, wait=True):
  logging.info("Updating " + table_name + " table read capacity to: " + str(read_capacity) + ", write capacity to: " + str(write_capacity))
  conn.update_table(table_name, {"ReadCapacityUnits": int(read_capacity), "WriteCapacityUnits": int(write_capacity)})

  # wait for provisioned throughput update completion
  if wait:
    wait_for_active_table(conn, table_name, "updated")

def do_backup(conn, table_name, read_capacity):
  logging.info("Starting backup for " + table_name + "..")

  # trash data, re-create subdir
  if os.path.exists(DUMP_PATH + "/" + table_name):
    shutil.rmtree(DUMP_PATH + "/" + table_name)
  mkdir_p(DUMP_PATH + "/" + table_name)

  # get table schema
  logging.info("Dumping table schema for " + table_name)
  f = open(DUMP_PATH + "/" + table_name + "/" + SCHEMA_FILE, "w+")
  table_desc = conn.describe_table(table_name)
  f.write(json.dumps(table_desc, indent=JSON_INDENT))
  f.close()

  # override table read capacity if specified
  if (read_capacity != None):
    original_read_capacity = table_desc["Table"]["ProvisionedThroughput"]["ReadCapacityUnits"]
    original_write_capacity = table_desc["Table"]["ProvisionedThroughput"]["WriteCapacityUnits"]
    update_provisioned_throughput(conn, table_name, read_capacity, original_write_capacity)

  # get table data
  logging.info("Dumping table items for " + table_name)
  mkdir_p(DUMP_PATH + "/" + table_name + "/" + DATA_DIR)

  i = 1
  last_evaluated_key = None

  while True:
    scanned_table = conn.scan(table_name, exclusive_start_key=last_evaluated_key)

    f = open(DUMP_PATH + "/" + table_name + "/" + DATA_DIR + "/" + str(i).zfill(4) + ".json", "w+")
    f.write(json.dumps(scanned_table, indent=JSON_INDENT))
    f.close()

    i += 1

    try:
      last_evaluated_key = scanned_table["LastEvaluatedKey"]
    except KeyError, e:
      break

  # revert back to original table read capacity if specified
  if (read_capacity != None):
    update_provisioned_throughput(conn, table_name, original_read_capacity, original_write_capacity, False)

  logging.info("Backup for " + table_name + " table completed. Time taken: " + str(datetime.datetime.now().replace(microsecond=0) - start_time))

def do_restore(conn, sleep_interval, source_table, destination_table, write_capacity):
  logging.info("Starting restore for " + source_table + " to " + destination_table + "..")

  # create table using schema
  table_data = json.load(open(DUMP_PATH + "/" + source_table + "/" + SCHEMA_FILE))
  table = table_data["Table"]
  table_attribute_definitions = table["AttributeDefinitions"]
  table_table_name = destination_table
  table_key_schema = table["KeySchema"]
  original_read_capacity = table["ProvisionedThroughput"]["ReadCapacityUnits"]
  original_write_capacity = table["ProvisionedThroughput"]["WriteCapacityUnits"]
  table_local_secondary_indexes = table.get("LocalSecondaryIndexes")
  table_global_secondary_indexes = table.get("GlobalSecondaryIndexes")

  # override table write capacity if specified, else use RESTORE_WRITE_CAPACITY
  if (write_capacity == None):
    write_capacity = RESTORE_WRITE_CAPACITY

  # temp provisioned throughput for restore
  table_provisioned_throughput = {"ReadCapacityUnits": int(original_read_capacity), "WriteCapacityUnits": int(write_capacity)}

  logging.info("Creating " + destination_table + " table with temp write capacity of " + str(write_capacity))

  while True:
    try:
      conn.create_table(table_attribute_definitions, table_table_name, table_key_schema, table_provisioned_throughput, table_local_secondary_indexes, table_global_secondary_indexes)
      break
    except boto.exception.JSONResponseError, e:
      if e.body["__type"] == "com.amazonaws.dynamodb.v20120810#LimitExceededException":
        logging.info("Limit exceeded, retrying creation of " + destination_table + "..")
        time.sleep(sleep_interval)
      else:
        logging.exception(e)
        sys.exit(1)

  # wait for table creation completion
  wait_for_active_table(conn, destination_table, "created")

  # read data files
  logging.info("Restoring data for " + destination_table + " table..")
  data_file_list = os.listdir(DUMP_PATH + "/" + source_table + "/" + DATA_DIR + "/")
  data_file_list.sort()

  items = []
  for data_file in data_file_list:
    item_data = json.load(open(DUMP_PATH + "/" + source_table + "/" + DATA_DIR + "/" + data_file))
    items.extend(item_data["Items"])

  # batch write data
  put_requests = []
  while len(items) > 0:
    put_requests.append({"PutRequest": {"Item": items.pop(0)}})

    # flush every MAX_BATCH_WRITE
    if len(put_requests) == MAX_BATCH_WRITE:
      logging.debug("Writing next " + str(MAX_BATCH_WRITE) + " items to " + destination_table + "..")
      batch_write(conn, sleep_interval, destination_table, put_requests)
      del put_requests[:]

  # flush remainder
  if len(put_requests) > 0:
    batch_write(conn, sleep_interval, destination_table, put_requests)

  # revert to original table write capacity
  update_provisioned_throughput(conn, destination_table, original_read_capacity, original_write_capacity, False)

  logging.info("Restore for " + source_table + " to " + destination_table + " table completed. Time taken: " + str(datetime.datetime.now().replace(microsecond=0) - start_time))

# parse args
parser = argparse.ArgumentParser(description="Simple DynamoDB backup/restore.")
parser.add_argument("-m", "--mode", help="'backup' or 'restore'")
parser.add_argument("-r", "--region", help="AWS region to use, e.g. 'us-west-1'. Use '" + LOCAL_REGION + "' for local DynamoDB testing.")
parser.add_argument("-s", "--srcTable", help="Source DynamoDB table name to backup or restore from, use 'tablename*' for wildcard prefix selection")
parser.add_argument("-d", "--destTable", help="Destination DynamoDB table name to backup or restore to, use 'tablename*' for wildcard prefix selection (uses '-' separator) [optional, defaults to source]")
parser.add_argument("--readCapacity", help="Change the temp read capacity of the DynamoDB table to backup from [optional]")
parser.add_argument("--writeCapacity", help="Change the temp write capacity of the DynamoDB table to restore to [defaults to " + str(RESTORE_WRITE_CAPACITY) + ", optional]")
parser.add_argument("--host", help="Host of local DynamoDB [required only for local]")
parser.add_argument("--port", help="Port of local DynamoDB [required only for local]")
parser.add_argument("--accessKey", help="Access key of local DynamoDB [required only for local]")
parser.add_argument("--secretKey", help="Secret key of local DynamoDB [required only for local]")
parser.add_argument("--log", help="Logging level - DEBUG|INFO|WARNING|ERROR|CRITICAL [optional]")
args = parser.parse_args()

# set log level
log_level = LOG_LEVEL
if args.log != None:
  log_level = args.log.upper()
logging.basicConfig(level=getattr(logging, log_level))

# instantiate connection
if args.region == LOCAL_REGION:
  conn = DynamoDBConnection(aws_access_key_id=args.accessKey, aws_secret_access_key=args.secretKey, host=args.host, port=int(args.port), is_secure=False)
  sleep_interval = LOCAL_SLEEP_INTERVAL
else:
  conn = boto.dynamodb2.connect_to_region(args.region, aws_access_key_id=args.accessKey, aws_secret_access_key=args.secretKey)
  sleep_interval = AWS_SLEEP_INTERVAL

# do backup/restore
start_time = datetime.datetime.now().replace(microsecond=0)
if args.mode == "backup":
  if args.srcTable.find("*") != -1:
    matching_backup_tables = get_table_name_matches(conn, args.srcTable)
    logging.info("Found " + str(len(matching_backup_tables)) + " table(s) in DynamoDB host to backup: " + ", ".join(matching_backup_tables))

    threads = []
    for table_name in matching_backup_tables:
      t = threading.Thread(target=do_backup, args=(conn, table_name, args.readCapacity,))
      threads.append(t)
      t.start()
      time.sleep(THREAD_START_DELAY)

    for thread in threads:
      thread.join()

    logging.info("Backup of table(s) " + args.srcTable + " completed!")
  else:
    do_backup(conn, args.srcTable, args.readCapacity)
elif args.mode == "restore":
  if args.destTable != None:
    dest_table = args.destTable
  else:
    dest_table = args.srcTable

  if dest_table.find("*") != -1:
    matching_destination_tables = get_table_name_matches(conn, dest_table)
    logging.info("Found " + str(len(matching_destination_tables)) + " table(s) in DynamoDB host to be deleted: " + ", ".join(matching_destination_tables))

    threads = []
    for table_name in matching_destination_tables:
      t = threading.Thread(target=delete_table, args=(conn, sleep_interval, table_name,))
      threads.append(t)
      t.start()
      time.sleep(THREAD_START_DELAY)

    for thread in threads:
      thread.join()

    matching_restore_tables = get_restore_table_matches(args.srcTable)
    logging.info("Found " + str(len(matching_restore_tables)) + " table(s) in " + DUMP_PATH + " to restore: " + ", ".join(matching_restore_tables))

    threads = []
    for source_table in matching_restore_tables:
      t = threading.Thread(target=do_restore, args=(conn, sleep_interval, source_table, change_prefix(source_table, args.srcTable, dest_table), args.writeCapacity,))
      threads.append(t)
      t.start()
      time.sleep(THREAD_START_DELAY)

    for thread in threads:
      thread.join()

    logging.info("Restore of table(s) " + args.srcTable + " to " +  dest_table + " completed!")
  else:
    delete_table(conn, sleep_interval, dest_table)
    do_restore(conn, sleep_interval, args.srcTable, dest_table, args.writeCapacity)
  <commit_msg>Added more conditional checks for edge cases for read/write capacities<commit_after>#!/usr/bin/env python
import boto.dynamodb2.layer1, json, sys, time, shutil, os, argparse, logging, datetime, threading
from boto.dynamodb2.layer1 import DynamoDBConnection

JSON_INDENT = 2
AWS_SLEEP_INTERVAL = 10 #seconds
LOCAL_SLEEP_INTERVAL = 1 #seconds
MAX_BATCH_WRITE = 25 #DynamoDB limit
SCHEMA_FILE = "schema.json"
DATA_DIR = "data"
MAX_RETRY = 3
LOCAL_REGION = "local"
LOG_LEVEL = "INFO"
DUMP_PATH = "dump"
RESTORE_WRITE_CAPACITY = 100
THREAD_START_DELAY = 1 #seconds

def get_table_name_matches(conn, table_name_wildcard):
  all_tables = []
  last_evaluated_table_name = None

  while True:
    table_list = conn.list_tables(exclusive_start_table_name=last_evaluated_table_name)
    all_tables.extend(table_list["TableNames"])

    try:
      last_evaluated_table_name = table_list["LastEvaluatedTableName"]
    except KeyError, e:
      break

  matching_tables = []
  for table_name in all_tables:
    if table_name.startswith(table_name_wildcard.split("*", 1)[0]):
      matching_tables.append(table_name)

  return matching_tables

def get_restore_table_matches(table_name_wildcard):
  matching_tables = []
  dir_list = os.listdir("./" + DUMP_PATH)
  for dir_name in dir_list:
    if dir_name.startswith(table_name_wildcard.split("*", 1)[0]):
      matching_tables.append(dir_name)

  return matching_tables

def change_prefix(source_table_name, source_wildcard, destination_wildcard):
  source_prefix = source_wildcard.split("*", 1)[0]
  destination_prefix = destination_wildcard.split("*", 1)[0]
  if source_table_name.split("-", 1)[0] == source_prefix:
    return destination_prefix + "-" + source_table_name.split("-", 1)[1]

def delete_table(conn, sleep_interval, table_name):
  # delete table if exists
  table_exist = True
  try:
    conn.delete_table(table_name)
  except boto.exception.JSONResponseError, e:
    if e.body["__type"] == "com.amazonaws.dynamodb.v20120810#ResourceNotFoundException":
      table_exist = False
      logging.info("Table does not exist in destination, deleting not necessary..")
      pass
    else:
      logging.exception(e)
      sys.exit(1)

  # if table exists, wait till deleted
  if table_exist:
    try:
      while True:
        logging.info("Waiting for " + table_name + " table to be deleted.. [" + conn.describe_table(table_name)["Table"]["TableStatus"] +"]")
        time.sleep(sleep_interval)
    except boto.exception.JSONResponseError, e:
      if e.body["__type"] == "com.amazonaws.dynamodb.v20120810#ResourceNotFoundException":
        logging.info(table_name + " table deleted.")
        pass
      else:
        logging.exception(e)
        sys.exit(1)

def mkdir_p(path):
  try:
    os.makedirs(path)
  except OSError as exc:
    if exc.errno == errno.EEXIST and os.path.isdir(path):
      pass
    else: raise

def batch_write(conn, sleep_interval, table_name, put_requests):
  request_items = {table_name: put_requests}
  i = 1
  while True:
    response = conn.batch_write_item(request_items)
    unprocessed_items = response["UnprocessedItems"]

    if len(unprocessed_items) == 0:
      break

    if len(unprocessed_items) > 0 and i <= MAX_RETRY:
      logging.info(str(len(unprocessed_items)) + " unprocessed items, retrying.. [" + str(i) + "]")
      request_items = unprocessed_items
      i += 1
      time.sleep(sleep_interval)
    else:
      logging.info("Max retries reached, failed to processed batch write: " + json.dumps(unprocessed_items, indent=JSON_INDENT))
      logging.info("Ignoring and continuing..")
      break

def wait_for_active_table(conn, table_name, verb):
  while True:
    if conn.describe_table(table_name)["Table"]["TableStatus"] != "ACTIVE":
      logging.info("Waiting for " + table_name + " table to be " + verb + ".. [" + conn.describe_table(table_name)["Table"]["TableStatus"] +"]")
      time.sleep(sleep_interval)
    else:
      logging.info(table_name + " " + verb + ".")
      break

def update_provisioned_throughput(conn, table_name, read_capacity, write_capacity, wait=True):
  logging.info("Updating " + table_name + " table read capacity to: " + str(read_capacity) + ", write capacity to: " + str(write_capacity))
  conn.update_table(table_name, {"ReadCapacityUnits": int(read_capacity), "WriteCapacityUnits": int(write_capacity)})

  # wait for provisioned throughput update completion
  if wait:
    wait_for_active_table(conn, table_name, "updated")

def do_backup(conn, table_name, read_capacity):
  logging.info("Starting backup for " + table_name + "..")

  # trash data, re-create subdir
  if os.path.exists(DUMP_PATH + "/" + table_name):
    shutil.rmtree(DUMP_PATH + "/" + table_name)
  mkdir_p(DUMP_PATH + "/" + table_name)

  # get table schema
  logging.info("Dumping table schema for " + table_name)
  f = open(DUMP_PATH + "/" + table_name + "/" + SCHEMA_FILE, "w+")
  table_desc = conn.describe_table(table_name)
  f.write(json.dumps(table_desc, indent=JSON_INDENT))
  f.close()

  original_read_capacity = table_desc["Table"]["ProvisionedThroughput"]["ReadCapacityUnits"]
  original_write_capacity = table_desc["Table"]["ProvisionedThroughput"]["WriteCapacityUnits"]

  # override table read capacity if specified
  if read_capacity != None and read_capacity != original_read_capacity:
    update_provisioned_throughput(conn, table_name, read_capacity, original_write_capacity)

  # get table data
  logging.info("Dumping table items for " + table_name)
  mkdir_p(DUMP_PATH + "/" + table_name + "/" + DATA_DIR)

  i = 1
  last_evaluated_key = None

  while True:
    scanned_table = conn.scan(table_name, exclusive_start_key=last_evaluated_key)

    f = open(DUMP_PATH + "/" + table_name + "/" + DATA_DIR + "/" + str(i).zfill(4) + ".json", "w+")
    f.write(json.dumps(scanned_table, indent=JSON_INDENT))
    f.close()

    i += 1

    try:
      last_evaluated_key = scanned_table["LastEvaluatedKey"]
    except KeyError, e:
      break

  # revert back to original table read capacity if specified
  if read_capacity != None and read_capacity != original_read_capacity:
    update_provisioned_throughput(conn, table_name, original_read_capacity, original_write_capacity, False)

  logging.info("Backup for " + table_name + " table completed. Time taken: " + str(datetime.datetime.now().replace(microsecond=0) - start_time))

def do_restore(conn, sleep_interval, source_table, destination_table, write_capacity):
  logging.info("Starting restore for " + source_table + " to " + destination_table + "..")

  # create table using schema
  table_data = json.load(open(DUMP_PATH + "/" + source_table + "/" + SCHEMA_FILE))
  table = table_data["Table"]
  table_attribute_definitions = table["AttributeDefinitions"]
  table_table_name = destination_table
  table_key_schema = table["KeySchema"]
  original_read_capacity = table["ProvisionedThroughput"]["ReadCapacityUnits"]
  original_write_capacity = table["ProvisionedThroughput"]["WriteCapacityUnits"]
  table_local_secondary_indexes = table.get("LocalSecondaryIndexes")
  table_global_secondary_indexes = table.get("GlobalSecondaryIndexes")

  # override table write capacity if specified, else use RESTORE_WRITE_CAPACITY if original write capacity is lower
  if write_capacity == None:
    if original_write_capacity < RESTORE_WRITE_CAPACITY:
      write_capacity = RESTORE_WRITE_CAPACITY
    else:
      write_capacity = original_write_capacity

  # temp provisioned throughput for restore
  table_provisioned_throughput = {"ReadCapacityUnits": int(original_read_capacity), "WriteCapacityUnits": int(write_capacity)}

  logging.info("Creating " + destination_table + " table with temp write capacity of " + str(write_capacity))

  while True:
    try:
      conn.create_table(table_attribute_definitions, table_table_name, table_key_schema, table_provisioned_throughput, table_local_secondary_indexes, table_global_secondary_indexes)
      break
    except boto.exception.JSONResponseError, e:
      if e.body["__type"] == "com.amazonaws.dynamodb.v20120810#LimitExceededException":
        logging.info("Limit exceeded, retrying creation of " + destination_table + "..")
        time.sleep(sleep_interval)
      else:
        logging.exception(e)
        sys.exit(1)

  # wait for table creation completion
  wait_for_active_table(conn, destination_table, "created")

  # read data files
  logging.info("Restoring data for " + destination_table + " table..")
  data_file_list = os.listdir(DUMP_PATH + "/" + source_table + "/" + DATA_DIR + "/")
  data_file_list.sort()

  items = []
  for data_file in data_file_list:
    item_data = json.load(open(DUMP_PATH + "/" + source_table + "/" + DATA_DIR + "/" + data_file))
    items.extend(item_data["Items"])

  # batch write data
  put_requests = []
  while len(items) > 0:
    put_requests.append({"PutRequest": {"Item": items.pop(0)}})

    # flush every MAX_BATCH_WRITE
    if len(put_requests) == MAX_BATCH_WRITE:
      logging.debug("Writing next " + str(MAX_BATCH_WRITE) + " items to " + destination_table + "..")
      batch_write(conn, sleep_interval, destination_table, put_requests)
      del put_requests[:]

  # flush remainder
  if len(put_requests) > 0:
    batch_write(conn, sleep_interval, destination_table, put_requests)

  # revert to original table write capacity if it has been modified
  if write_capacity != original_write_capacity:
    update_provisioned_throughput(conn, destination_table, original_read_capacity, original_write_capacity, False)

  logging.info("Restore for " + source_table + " to " + destination_table + " table completed. Time taken: " + str(datetime.datetime.now().replace(microsecond=0) - start_time))

# parse args
parser = argparse.ArgumentParser(description="Simple DynamoDB backup/restore.")
parser.add_argument("-m", "--mode", help="'backup' or 'restore'")
parser.add_argument("-r", "--region", help="AWS region to use, e.g. 'us-west-1'. Use '" + LOCAL_REGION + "' for local DynamoDB testing.")
parser.add_argument("-s", "--srcTable", help="Source DynamoDB table name to backup or restore from, use 'tablename*' for wildcard prefix selection")
parser.add_argument("-d", "--destTable", help="Destination DynamoDB table name to backup or restore to, use 'tablename*' for wildcard prefix selection (uses '-' separator) [optional, defaults to source]")
parser.add_argument("--readCapacity", help="Change the temp read capacity of the DynamoDB table to backup from [optional]")
parser.add_argument("--writeCapacity", help="Change the temp write capacity of the DynamoDB table to restore to [defaults to " + str(RESTORE_WRITE_CAPACITY) + ", optional]")
parser.add_argument("--host", help="Host of local DynamoDB [required only for local]")
parser.add_argument("--port", help="Port of local DynamoDB [required only for local]")
parser.add_argument("--accessKey", help="Access key of local DynamoDB [required only for local]")
parser.add_argument("--secretKey", help="Secret key of local DynamoDB [required only for local]")
parser.add_argument("--log", help="Logging level - DEBUG|INFO|WARNING|ERROR|CRITICAL [optional]")
args = parser.parse_args()

# set log level
log_level = LOG_LEVEL
if args.log != None:
  log_level = args.log.upper()
logging.basicConfig(level=getattr(logging, log_level))

# instantiate connection
if args.region == LOCAL_REGION:
  conn = DynamoDBConnection(aws_access_key_id=args.accessKey, aws_secret_access_key=args.secretKey, host=args.host, port=int(args.port), is_secure=False)
  sleep_interval = LOCAL_SLEEP_INTERVAL
else:
  conn = boto.dynamodb2.connect_to_region(args.region, aws_access_key_id=args.accessKey, aws_secret_access_key=args.secretKey)
  sleep_interval = AWS_SLEEP_INTERVAL

# do backup/restore
start_time = datetime.datetime.now().replace(microsecond=0)
if args.mode == "backup":
  if args.srcTable.find("*") != -1:
    matching_backup_tables = get_table_name_matches(conn, args.srcTable)
    logging.info("Found " + str(len(matching_backup_tables)) + " table(s) in DynamoDB host to backup: " + ", ".join(matching_backup_tables))

    threads = []
    for table_name in matching_backup_tables:
      t = threading.Thread(target=do_backup, args=(conn, table_name, args.readCapacity,))
      threads.append(t)
      t.start()
      time.sleep(THREAD_START_DELAY)

    for thread in threads:
      thread.join()

    logging.info("Backup of table(s) " + args.srcTable + " completed!")
  else:
    do_backup(conn, args.srcTable, args.readCapacity)
elif args.mode == "restore":
  if args.destTable != None:
    dest_table = args.destTable
  else:
    dest_table = args.srcTable

  if dest_table.find("*") != -1:
    matching_destination_tables = get_table_name_matches(conn, dest_table)
    logging.info("Found " + str(len(matching_destination_tables)) + " table(s) in DynamoDB host to be deleted: " + ", ".join(matching_destination_tables))

    threads = []
    for table_name in matching_destination_tables:
      t = threading.Thread(target=delete_table, args=(conn, sleep_interval, table_name,))
      threads.append(t)
      t.start()
      time.sleep(THREAD_START_DELAY)

    for thread in threads:
      thread.join()

    matching_restore_tables = get_restore_table_matches(args.srcTable)
    logging.info("Found " + str(len(matching_restore_tables)) + " table(s) in " + DUMP_PATH + " to restore: " + ", ".join(matching_restore_tables))

    threads = []
    for source_table in matching_restore_tables:
      t = threading.Thread(target=do_restore, args=(conn, sleep_interval, source_table, change_prefix(source_table, args.srcTable, dest_table), args.writeCapacity,))
      threads.append(t)
      t.start()
      time.sleep(THREAD_START_DELAY)

    for thread in threads:
      thread.join()

    logging.info("Restore of table(s) " + args.srcTable + " to " +  dest_table + " completed!")
  else:
    delete_table(conn, sleep_interval, dest_table)
    do_restore(conn, sleep_interval, args.srcTable, dest_table, args.writeCapacity)
  <|endoftext|>


<commit_before>import os
from flask import Flask
from flask import render_template
from flask import request, redirect, url_for
from werkzeug import secure_filename
import hashlib

from pymongo import Connection

import dpkt
import sys
import socket

from celery import Celery

from requests import Request, Session

from vt import vtAPI as vt

# FIXME: move to config.py
ALLOWED_EXTENSIONS = set(['pcap'])

def create_app():
    return Flask("ekanalyzer")

app = create_app()
app.config.from_pyfile('config.py')

connection = Connection(app.config['MONGODB_SERVER'] , app.config['MONGODB_PORT'])
db = connection.ekanalyzer

app.debug = True

celery = Celery('ekanalyzer', broker=app.config['BROKER_URL'] )

@celery.task
def perform_results(hash):
    try:

        pcap = {'hash' : hash}

        result = db.pcap.find(pcap)

        if result.count() > 0:
            return
        else:
            db.pcap.insert(pcap)

        f = open(app.config['UPLOAD_FOLDER'] + hash)
        pcap = dpkt.pcap.Reader(f)
        for ts, buf in pcap:
            eth = dpkt.ethernet.Ethernet(buf)
            ip = eth.data
            tcp = ip.data
            # FIXME: assuming only http traffic on port 80
            if tcp.dport == 80 and len(tcp.data) > 0:
                http = dpkt.http.Request(tcp.data)
                ipaddress =  socket.inet_ntoa(ip.dst)

                data = { 'ip' : ipaddress,
                         'uri' : http.uri,
                         'method' : http.method,
                         'data' : http.data,
                         'headers' : http.headers,
                         'hash': hash
                       }

                db.requests.insert(data)

        print "Data imported"
        status = process_requests(hash)

    except NameError as e:
        print e
    except :
        print "Unexpected error:", sys.exc_info()
        pass

def process_requests(hash):
    request = { 'hash' : hash}
    result = db.requests.find(request)
    for r in result:
        print process_request.delay(r['ip'], r['uri'], r['method'], r['headers'], r['data'], hash)

@celery.task
def process_request(ip, uri, method, headers, data, hash):

    user_agents = app.config['USER_AGENTS']
    for user_agent in user_agents:
        headers['user-agent'] = user_agent

        #FIXME: port 80
        url = "http://{0}:80{1}".format(ip, uri)

        s = Session()
        req = Request(method, url,
            data=data,
            headers=headers
        )
        prepped = req.prepare()

        resp = s.send(prepped)

        m = hashlib.md5()
        m.update(user_agent)
        UA = m.hexdigest()

        fpath = "workspace/" + hash + "/" + UA + "/" + headers['host'] + uri
        dpath = os.path.dirname(fpath)

        if not os.path.exists(dpath):
            os.makedirs(dpath)

        with open(fpath, "w") as f:
            f.write(resp.content)

        m = hashlib.sha256()
        m.update(resp.content)
        response = m.hexdigest()

        analysis_data = { 'hash': hash, 'user-agent': user_agent, 'host': headers['host'], 'uri' : uri, 'data' : data, 'status_code': resp.status_code, 'response': response }

        db.analysis.insert(analysis_data)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS

@app.route('/upload-ek/', methods=['POST'])
def upload_file():
    file = request.files['pcap']
    if file and allowed_file(file.filename):

        hash = hashlib.sha256()

        try:
            # FIXME: it should be saved before calculate sha256
            for chunk in file.chunks():
                hash.update(chunk)
        except:
            print "Unexpected error:", sys.exc_info()
        finally:
            file.seek(0)
            hash_name = "%s" % (hash.hexdigest())
            file.save(os.path.join(app.config['UPLOAD_FOLDER'], hash_name))
            return redirect(url_for('launch', hash=hash_name))

@app.route('/launch/<hash>/')
def launch(hash):
    perform_results.delay(hash)
    return render_template('launch.html', hash=hash)

@app.route('/view/<hash>/')
def view(hash):
    return render_template('view.html', hash=hash)

@app.route('/')
def index():
    return render_template('index.html')

if __name__ == "__main__":
    app.run(debug=True)

<commit_msg>Some refactoring<commit_after>import os
from flask import Flask
from flask import render_template
from flask import request, redirect, url_for
from werkzeug import secure_filename
import hashlib

from pymongo import Connection

import dpkt
import sys
import socket

from celery import Celery

from requests import Request, Session

from vt import vtAPI as vt

# FIXME: move to config.py
ALLOWED_EXTENSIONS = set(['pcap'])

def create_app():
    return Flask("ekanalyzer")

app = create_app()
app.config.from_pyfile('config.py')

connection = Connection(app.config['MONGODB_SERVER'] , app.config['MONGODB_PORT'])
db = connection.ekanalyzer

app.debug = True

celery = Celery('ekanalyzer', broker=app.config['BROKER_URL'] )

@celery.task
def perform_results(hash):
    try:

        pcap = {'id' : hash}

        result = db.pcap.find(pcap)

        if result.count() > 0:
            print "Ya existe pcap no se hace"
            return
        else:
            print "Vamos a hacer el pcap"
            db.pcap.insert(pcap)

        f = open(app.config['UPLOAD_FOLDER'] + hash)
        pcap = dpkt.pcap.Reader(f)
        for ts, buf in pcap:
            eth = dpkt.ethernet.Ethernet(buf)
            ip = eth.data
            tcp = ip.data
            # FIXME: assuming only http traffic on port 80
            if tcp.dport == 80 and len(tcp.data) > 0:
                http = dpkt.http.Request(tcp.data)
                ipaddress =  socket.inet_ntoa(ip.dst)

                data = { 'ip' : ipaddress,
                         'uri' : http.uri,
                         'method' : http.method,
                         'data' : http.data,
                         'headers' : http.headers,
                         'id': hash
                       }

                db.requests.insert(data)

        print "Data imported"
        status = process_requests(hash)

    except NameError as e:
        print e
    except :
        print "Unexpected error:", sys.exc_info()
        pass

def process_requests(id):
    request = { 'id' : id}
    result = db.requests.find(request)
    for r in result:
        print process_request.delay(r['ip'], r['uri'], r['method'], r['headers'], r['data'], id)

@celery.task
def process_request(ip, uri, method, headers, data, id):

    user_agents = app.config['USER_AGENTS']
    for user_agent in user_agents:
        headers['user-agent'] = user_agent

        #FIXME: port 80
        url = "http://{0}:80{1}".format(ip, uri)

        s = Session()
        req = Request(method, url,
            data=data,
            headers=headers
        )
        prepped = req.prepare()

        resp = s.send(prepped)

        #user agent hash
        m = hashlib.md5()
        m.update(user_agent)
        UA = m.hexdigest()

        fpath = "workspace/" + id + "/" + UA + "/" + headers['host'] + uri
        dpath = os.path.dirname(fpath)

        if not os.path.exists(dpath):
            os.makedirs(dpath)

        response = resp.content

        with open(fpath, "w") as f:
            f.write(response)

        m.update(resp.content)
        response = m.hexdigest()

        # response
        m = hashlib.sha256()
        m.update(response)
        hash = m.hexdigest()

        analysis_data = { 'id': id, 'user-agent': user_agent, 'host': headers['host'], 'uri' : uri, 'data' : data, 'status_code': resp.status_code, 'hash': hash  }

        db.analysis.insert(analysis_data)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS

@app.route('/upload-ek/', methods=['POST'])
def upload_file():
    file = request.files['pcap']
    if file and allowed_file(file.filename):

        hash = hashlib.sha256()

        try:
            # FIXME: it should be saved before calculate sha256
            hash.update(file.read())
        except:
            print "Unexpected error:", sys.exc_info()
        finally:
            file.seek(0)
            hash_name = "%s" % (hash.hexdigest())
            file.save(os.path.join(app.config['UPLOAD_FOLDER'], hash_name))
            return redirect(url_for('launch', hash=hash_name))

@app.route('/launch/<hash>/')
def launch(hash):
    perform_results.delay(hash)
    return render_template('launch.html', hash=hash)

@app.route('/view/<hash>/')
def view(hash):
    return render_template('view.html', hash=hash)

@app.route('/')
def index():
    return render_template('index.html')

if __name__ == "__main__":
    app.run(debug=True)

<|endoftext|>


<commit_before>import os
from flask import Flask
from flask import render_template
from flask import request, redirect, url_for
from werkzeug import secure_filename
import hashlib

from pymongo import Connection
from bson.code import Code

import dpkt
import sys
import socket

from celery import Celery

import requests
from requests import Request, Session

import magic
import zlib

import yara
import pyclamd

# FIXME: move to config.py
ALLOWED_EXTENSIONS = set(['pcap'])
rules = yara.compile(filepath='yara/ekanalyzer.yar')

cd = pyclamd.ClamdAgnostic()

def create_app():
    return Flask("ekanalyzer")

app = create_app()
app.config.from_pyfile('config.py')

connection = Connection(app.config['MONGODB_SERVER'] , app.config['MONGODB_PORT'])
db = connection.ekanalyzer

app.debug = True

celery = Celery('ekanalyzer', broker=app.config['BROKER_URL'] )

@celery.task
def perform_results(hash):
    try:

        pcap = {'id' : hash}

        result = db.pcap.find(pcap)

        if result.count() > 0:
            return
        else:
            db.pcap.insert(pcap)

        f = open(app.config['UPLOAD_FOLDER'] + hash, "rb")
        pcap = dpkt.pcap.Reader(f)
        for ts, buf in pcap:
            eth = dpkt.ethernet.Ethernet(buf)
            ip = eth.data
            tcp = ip.data
            # FIXME: assuming only http traffic on port 80
            if tcp.dport == 80 and len(tcp.data) > 0:
                http = dpkt.http.Request(tcp.data)
                ipaddress =  socket.inet_ntoa(ip.dst)

                data = { 'ip' : ipaddress,
                         'uri' : http.uri,
                         'method' : http.method,
                         'data' : http.data,
                         'headers' : http.headers,
                         'id': hash
                       }
                db.requests.insert(data)
            #else:
            #    print "Port is " + str(tcp.dport)

    except dpkt.NeedData as e:
        print e
    except AttributeError as e:
        print e
    except NameError as e:
        print e
    except :
        print "Unexpected error:", sys.exc_info()
    finally:
        status = process_requests(hash)

def process_requests(id):
    request = { 'id' : id}
    result = db.requests.find(request)
    for r in result:
        print process_request.delay(r['ip'], r['uri'], r['method'], r['headers'], r['data'], id)

@celery.task
def process_request(ip, uri, method, headers, data, id):

    user_agents = app.config['USER_AGENTS']

    # FIXME: check case
    if 'user-agent' in headers:
        user_agents.append(headers['user-agent'])
    else:
        user_agents.append("")

    for user_agent in user_agents:
        headers['user-agent'] = user_agent

        #FIXME: port 80
        #FIXME: ConnectionError
        url = "http://{0}:80{1}".format(ip, uri)

        #proxies = {
        # "http": "http://127.0.0.1:8080"
        #}

        s = Session()
        req = Request(method, url,
            data=data,
            headers=headers
        )
        prepped = req.prepare()

        resp = s.send(prepped,
            #proxies=proxies
        )

        #user agent hash
        m = hashlib.md5()
        m.update(user_agent)
        UA = m.hexdigest()

        fpath = "workspace/" + id + "/" + UA + "/" + headers['host'] + uri
        dpath = os.path.dirname(fpath)

        if not os.path.exists(dpath):
            os.makedirs(dpath)

        response = resp.content

        # FIXME: uris ending with / are not saved properly
        if not os.path.isdir(fpath):
            with open(fpath, "w") as f:
                f.write(response)

        # response hash
        m = hashlib.sha256()
        m.update(response)
        hash = m.hexdigest()

        # filetype & mimetype
        filetype = magic.from_buffer(response)
        mimetype = magic.from_buffer(response, mime=True)

        vt_report = None

        tags = { 'clean' : 0, 'suspicious' : 0, 'malicious' : 0 }
        malicious = False

        ymatches = None

        #
        # This function uses response (buffer) and fpath (path to file)
        # FIX this as soon as the "/" bug be fixed
        #

        # Send to VT
        if mimetype == "application/octet-stream" \
            or mimetype == "application/java-archive" \
            or mimetype == "application/zip" \
            or mimetype == "'application/pdf'" \
            or mimetype == "application/x-shockwave-flash":

            # FIME: cache for repeated analysis
            parameters = {"resource": hash, "apikey": app.config["VIRUSTOTAL_API_KEY"]}
            r = requests.post('https://www.virustotal.com/vtapi/v2/file/report', params=parameters)
            try:
                vt_report = r.json()
                if vt_report['positives'] > 0:
                    tags['malicious'] += 1
                    malicious = True
            except:
                #print "Unexpected error:", sys.exc_info()
                vt_report = None

        # Prepare for YARA
        # FIXME: ZWS http://malware-traffic-analysis.net/2014/09/23/index.html
        if mimetype == "application/x-shockwave-flash" and filetype.find("CWS"):
            #print "compressed SWF detected"
            f = open(fpath, "rb")
            f.read(3) # skip 3 bytes
            tmp = 'FWS' + f.read(5) + zlib.decompress(f.read())
            decompressed = fpath + ".decompressed"
            with open(decompressed, "w") as f:
                f.write(tmp)
            unpacked =tmp
        else:
            unpacked = response

        ymatches = rules.match(data=unpacked)
        if not bool(ymatches):
            ymatches = None
        else:
            tags['suspicious'] += 1
            malicious = True

        # ClamAV analysis
        clamav = cd.scan_stream(unpacked)
        if clamav:
            malicious = True
            tags['malicious'] += 1

        #FIXME: add html/javascript analysis here

        #FIXME: add peepdf based analysis here

        # Review tags before analysis
        if tags['malicious'] == 0 and tags['suspicious'] == 0:
            tags['clean'] = 1

        analysis_data = { 'id': id, 'tags': tags, 'malicious': malicious, 'filetype': filetype,'mimetype': mimetype, 'yara' : ymatches, 'clamav' : clamav, 'user-agent': user_agent, 'UA' : UA,  'host': headers['host'], 'uri' : uri, 'data' : data, 'status_code': resp.status_code, 'hash': hash , 'vt' : vt_report }

        db.analysis.insert(analysis_data)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS

@app.route('/upload-ek/', methods=['POST'])
def upload_file():
    file = request.files['pcap']

    if file and allowed_file(file.filename):

        hash = hashlib.sha256()

        try:
            # FIXME: it should be saved before calculate sha256
            hash.update(file.read())
        except:
            print "Unexpected error:", sys.exc_info()
        finally:
            file.seek(0)
            hash_name = "%s" % (hash.hexdigest())
            file.save(os.path.join(app.config['UPLOAD_FOLDER'], hash_name))
            return redirect(url_for('launch', hash=hash_name))

@app.route('/launch/<hash>/')
def launch(hash):
    perform_results.delay(hash)
    return render_template('launch.html', hash=hash)

@app.route('/view/<hash>/')
def view(hash):

    h = { "_id.hash" : hash }

    #requests = db.analysis.find(h)

    # FIXME: this map/reduce is executed each time view is requested
    map = Code("function () {"
        "  emit({ hash : this['id'], UA : this.UA, 'user-agent' : this['user-agent']}, {malicious: this.tags.malicious, clean: this.tags.clean, suspicious:this.tags.suspicious});"
        "}")

    reduce = Code("function (key, vals) {"
        "  var result = {malicious:0, suspicious:0, clean:0 };"
        "  vals.forEach(function (value) {result.malicious += value.malicious; result.clean += value.clean; result.suspicious += value.suspicious; });"
        "  return result;"
        "}")

    results = db.analysis.map_reduce(map, reduce, 'malicious')

    found = results.find(h)
    requests = []

    for i in found:
        print i
        requests.append(i)

    original_request = db.requests.find_one({"id": hash})

    original_ua = ''

    if original_request:
        original_ua = original_request['headers']['user-agent']

    return render_template('view.html', requests=requests, original_ua=original_ua)

@app.route('/list')
def list():

    pcaps = db.pcap.find()

    analysis = []

    malicious = False

    for pcap in pcaps:
        h = { 'id' : pcap['id']}
        queries = db.analysis.find(h)
        details = []
        tags = { 'malicious' : 0, 'suspicious': 0, 'clean': 0}
        for query in queries:
            #print query
            if query['tags']['malicious']:
               tags['malicious'] += 1
            if query['tags']['suspicious']:
               tags['suspicious'] += 1
            if query['tags']['clean']:
               tags['clean'] += 1

        analysis.append( {pcap['id'] : tags})
    return render_template('list.html', analysis=analysis)

@app.route('/details/<hash>/<ua>/')
def details(hash, ua):
    user_agent = { 'UA' : ua, 'id' : hash}
    requests = db.analysis.find(user_agent)

    return render_template('details.html', requests=requests)

@app.route('/')
def index():
    return render_template('index.html')

if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0")

<commit_msg>Minor changes<commit_after>import os
from flask import Flask
from flask import render_template
from flask import request, redirect, url_for
from werkzeug import secure_filename
import hashlib

from pymongo import Connection
from bson.code import Code

import dpkt
import sys
import socket

from celery import Celery

import requests
from requests import Request, Session

import magic
import zlib

import yara
import pyclamd

# FIXME: move to config.py
ALLOWED_EXTENSIONS = set(['pcap'])
rules = yara.compile(filepath='yara/ekanalyzer.yar')

cd = pyclamd.ClamdAgnostic()

def create_app():
    return Flask("ekanalyzer")

app = create_app()
app.config.from_pyfile('config.py')

connection = Connection(app.config['MONGODB_SERVER'] , app.config['MONGODB_PORT'])
db = connection.ekanalyzer

app.debug = True

celery = Celery('ekanalyzer', broker=app.config['BROKER_URL'] )

@celery.task
def perform_results(hash):
    try:

        pcap = {'id' : hash}

        result = db.pcap.find(pcap)

        if result.count() > 0:
            return
        else:
            db.pcap.insert(pcap)

        f = open(app.config['UPLOAD_FOLDER'] + hash, "rb")
        pcap = dpkt.pcap.Reader(f)
        for ts, buf in pcap:
            eth = dpkt.ethernet.Ethernet(buf)
            ip = eth.data
            tcp = ip.data
            # FIXME: assuming only http traffic on port 80
            if tcp.dport == 80 and len(tcp.data) > 0:
                http = dpkt.http.Request(tcp.data)
                ipaddress =  socket.inet_ntoa(ip.dst)

                data = { 'ip' : ipaddress,
                         'uri' : http.uri,
                         'method' : http.method,
                         'data' : http.data,
                         'headers' : http.headers,
                         'id': hash
                       }
                db.requests.insert(data)
            #else:
            #    print "Port is " + str(tcp.dport)

    except dpkt.NeedData as e:
        print e
    except AttributeError as e:
        print e
    except NameError as e:
        print e
    except :
        print "Unexpected error:", sys.exc_info()
    finally:
        status = process_requests(hash)

def process_requests(id):
    request = { 'id' : id}
    result = db.requests.find(request)
    for r in result:
        print process_request.delay(r['ip'], r['uri'], r['method'], r['headers'], r['data'], id)

@celery.task
def process_request(ip, uri, method, headers, data, id):

    user_agents = app.config['USER_AGENTS']

    # FIXME: check case
    if 'user-agent' in headers:
        user_agents.append(headers['user-agent'])
    else:
        user_agents.append("")

    for user_agent in user_agents:
        headers['user-agent'] = user_agent

        #FIXME: port 80
        #FIXME: ConnectionError
        url = "http://{0}:80{1}".format(ip, uri)

        #proxies = {
        # "http": "http://127.0.0.1:8080"
        #}

        s = Session()
        req = Request(method, url,
            data=data,
            headers=headers
        )
        prepped = req.prepare()

        resp = s.send(prepped,
            #proxies=proxies
        )

        #user agent hash
        m = hashlib.md5()
        m.update(user_agent)
        UA = m.hexdigest()

        fpath = "workspace/" + id + "/" + UA + "/" + headers['host'] + uri
        dpath = os.path.dirname(fpath)

        if not os.path.exists(dpath):
            os.makedirs(dpath)

        response = resp.content

        # FIXME: uris ending with / are not saved properly
        if not os.path.isdir(fpath):
            with open(fpath, "w") as f:
                f.write(response)

        # response hash
        m = hashlib.sha256()
        m.update(response)
        hash = m.hexdigest()

        # filetype & mimetype
        filetype = magic.from_buffer(response)
        mimetype = magic.from_buffer(response, mime=True)

        vt_report = None

        tags = { 'clean' : 0, 'suspicious' : 0, 'malicious' : 0 }
        malicious = False

        ymatches = None

        #
        # This function uses response (buffer) and fpath (path to file)
        # FIX this as soon as the "/" bug be fixed
        #

        # Send to VT
        if mimetype == "application/octet-stream" \
            or mimetype == "application/java-archive" \
            or mimetype == "application/zip" \
            or mimetype == "'application/pdf'" \
            or mimetype == "application/x-shockwave-flash":

            # FIME: cache for repeated analysis
            parameters = {"resource": hash, "apikey": app.config["VIRUSTOTAL_API_KEY"]}
            r = requests.post('https://www.virustotal.com/vtapi/v2/file/report', params=parameters)
            try:
                vt_report = r.json()
                if vt_report['positives'] > 0:
                    tags['malicious'] += 1
                    malicious = True
            except:
                #print "Unexpected error:", sys.exc_info()
                vt_report = None

        # Prepare for YARA
        # FIXME: ZWS http://malware-traffic-analysis.net/2014/09/23/index.html
        if mimetype == "application/x-shockwave-flash" and filetype.find("CWS"):
            #print "compressed SWF detected"
            f = open(fpath, "rb")
            f.read(3) # skip 3 bytes
            tmp = 'FWS' + f.read(5) + zlib.decompress(f.read())
            decompressed = fpath + ".decompressed"
            with open(decompressed, "w") as f:
                f.write(tmp)
            unpacked =tmp
        else:
            unpacked = response

        ymatches = rules.match(data=unpacked)
        if not bool(ymatches):
            ymatches = None
        else:
            tags['suspicious'] += 1
            malicious = True

        # ClamAV analysis
        clamav = cd.scan_stream(unpacked)
        if clamav:
            malicious = True
            tags['malicious'] += 1

        #FIXME: add html/javascript analysis here

        #FIXME: add peepdf based analysis here

        # Review tags before analysis
        if tags['malicious'] == 0 and tags['suspicious'] == 0:
            tags['clean'] = 1

        analysis_data = { 'id': id, 'tags': tags, 'malicious': malicious, 'filetype': filetype,'mimetype': mimetype, 'yara' : ymatches, 'clamav' : clamav, 'user-agent': user_agent, 'UA' : UA,  'host': headers['host'], 'uri' : uri, 'data' : data, 'status_code': resp.status_code, 'hash': hash , 'vt' : vt_report }

        db.analysis.insert(analysis_data)

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS

@app.route('/upload-ek/', methods=['POST'])
def upload_file():
    file = request.files['pcap']

    if file and allowed_file(file.filename):

        hash = hashlib.sha256()

        try:
            # FIXME: it should be saved before calculate sha256
            hash.update(file.read())
        except:
            print "Unexpected error:", sys.exc_info()
        finally:
            file.seek(0)
            hash_name = "%s" % (hash.hexdigest())
            file.save(os.path.join(app.config['UPLOAD_FOLDER'], hash_name))
            return redirect(url_for('launch', hash=hash_name))

@app.route('/launch/<hash>/')
def launch(hash):
    perform_results.delay(hash)
    return render_template('launch.html', hash=hash)

@app.route('/view/<hash>/')
def view(hash):

    h = { "_id.hash" : hash }

    #requests = db.analysis.find(h)

    # FIXME: this map/reduce is executed each time view is requested
    map = Code("function () {"
        "  emit({ hash : this['id'], UA : this.UA, 'user-agent' : this['user-agent']}, {malicious: this.tags.malicious, clean: this.tags.clean, suspicious:this.tags.suspicious});"
        "}")

    reduce = Code("function (key, vals) {"
        "  var result = {malicious:0, suspicious:0, clean:0 };"
        "  vals.forEach(function (value) {result.malicious += value.malicious; result.clean += value.clean; result.suspicious += value.suspicious; });"
        "  return result;"
        "}")

    results = db.analysis.map_reduce(map, reduce, 'malicious')

    found = results.find(h)
    requests = []

    for i in found:
        #print i
        requests.append(i)

    original_request = db.requests.find_one({"id": hash})

    original_ua = ''

    if original_request:
        original_ua = original_request['headers']['user-agent']

    return render_template('view.html', requests=requests, original_ua=original_ua)

@app.route('/list')
def list():

    pcaps = db.pcap.find()

    analysis = []

    malicious = False

    for pcap in pcaps:
        h = { 'id' : pcap['id']}
        queries = db.analysis.find(h)
        details = []
        tags = { 'malicious' : 0, 'suspicious': 0, 'clean': 0}
        for query in queries:
            #print query
            if query['tags']['malicious']:
               tags['malicious'] += 1
            if query['tags']['suspicious']:
               tags['suspicious'] += 1
            if query['tags']['clean']:
               tags['clean'] += 1

        analysis.append( {pcap['id'] : tags})
    return render_template('list.html', analysis=analysis)

@app.route('/details/<hash>/<ua>/')
def details(hash, ua):
    user_agent = { 'UA' : ua, 'id' : hash}
    requests = db.analysis.find(user_agent)

    return render_template('details.html', requests=requests)

@app.route('/')
def index():
    return render_template('index.html')

if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0")

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import re
from utils import detect_string, gen_hashcode

IOWAIT_THRESHOLD = 35.0
TOTAL_THRESHOLD = 40.0

def detect_keyword(exp, content, mode=0):
    match = re.findall(exp, content)
    if mode == 1:
        return match if match else []
    return match[0] if match else None

def java_binder(cpu_usage, mainstack, process):
    detail = detect_keyword(r'\s+at\s+android\.os\.BinderProxy\.transact\(Native Method\)\n\s+(at\s+.+)\n', mainstack)
    if detail:
        return process, detail, None
    return None, None, None

def system_busy(cpu_usage, mainstack, process):
    total = detect_keyword(r'(\d+(?:\.\d+)?)\%\s+TOTAL', cpu_usage)
    iowait = detect_keyword(r'(\d+(?:\.\d+)?)\%\s+iowait', cpu_usage)
    if iowait and float(iowait) >= IOWAIT_THRESHOLD:
        return "iowait", "CPU usage: iowait >= 35%", None
    if total and float(total) >= TOTAL_THRESHOLD:
        topprocess = detect_keyword(r'\s+\d+(?:\.\d+)?\%\s+(?:\d+)?/([\w-]+(?:\.[\w-]+)*)', cpu_usage)
        return topprocess, ("CPU usage: TOTAL >= 40%, " + topprocess + " occupied the most."), None
    return None, None, None

def native_binder(cpu_usage, mainstack, process):
    details = detect_keyword(r'\s+((?:/.+)+(?:\s+\(.+\))?)', mainstack, 1)
    binder = False
    for detail in details:
        if not binder and 'libbinder.so' in detail:
            binder = True
        if binder and 'libbinder.so' not in detail:
            for key in ['libmedia\.so', 'libtvmanager\.so', 'libaudiomanager\.so', 'libgui\.so', 'libsystem_server\.so', 'libchannelmanager\.so']:
                if detect_keyword((r'%s\s+\(.+\)' %(key)), detail):
                    return key.replace("\\", ""), detail, None
    return None, None, None

def native_lock(cpu_usage, mainstack, process):
    details = detect_keyword(r'\s+((?:/.+)+(?:\s+\(.+\))?)', mainstack, 1)
    lock = False
    for detail in details:
        if not lock and 'libc.so (__futex_syscall' in detail:
            lock = True
        if lock and 'libc.so' not in detail:
            return process, detail, None
    return None, None, None

def java_lock(cpu_usage, mainstack, process):
    for exp in [r'\s+-\s+(waiting\s+on\s+<.+>\s+(?:\(a\s+\w+(?:[.$]\w+)+\))?)', r'\s+-\s+(waiting\s+to\s+lock\s+<.+>\s+(?:\(a\s+\w+(?:[.$]\w+)+\))?)']:
        detail = detect_keyword(exp, mainstack)
        if detail:
            return process, re.sub(r'<\w+>', 'addr.', detail), None
    return None, None, None

def mthread_busy(cpu_usage, mainstack, process):
    for exp in [r'\s+(at\s+java\.io\..+)', r'\s+(at\s+android\.database\.sqlite\..+)']:
        detail = detect_keyword(exp, mainstack)
        if detail:
            return process, detail, None
    return None, None, None

def auto_recover(cpu_usage, mainstack, process):
    details = detect_keyword(r'\s+(at\s+\w+(?:[.$]\w+)+(?:\(.+\))?)', mainstack, 1)
    if details and "android.os.MessageQueue.nativePollOnce" in details[0]:
        return None, None, "RECOVER"
    return None, None, None

def detect_basic_info(logcat):
    process, cpu_usage, mainstack = None, None, None
    for minor_part in logcat.split('\n\n'):
        if not process:
            process = detect_string(minor_part, r'^Process:\s+(\w+(?:[.$_]\w+)+)')
        if not cpu_usage and detect_string(minor_part, r'(CPU\s+usage\s+from\s+-?\w+\s+to\s+-?\w+)'):
            cpu_usage = minor_part
        if not mainstack and detect_string(minor_part, r'^(\"main\")\s+'):
            mainstack = minor_part
    return process, cpu_usage, mainstack

METHODS = [java_binder, system_busy, native_binder, native_lock, java_lock, mthread_busy, auto_recover]

def app_anr(logcat, headers):
    process, cpu_usage, mainstack = detect_basic_info(logcat)
    if process and cpu_usage and mainstack:
        for method in METHODS:
            process, detail, tag = method(cpu_usage, mainstack, process)
            if tag == "RECOVER":
                break
            if process and detail:
                md5 = gen_hashcode({'issue_owner': process, 'detail': detail})
                return md5, {'issue_owner': process, 'detail': detail}, None
    return None, None, None
<commit_msg>fix anr error<commit_after># -*- coding: utf-8 -*-

import re
from utils import detect_string, gen_hashcode

IOWAIT_THRESHOLD = 35.0
TOTAL_THRESHOLD = 40.0

def detect_keyword(exp, content, mode=0):
    match = re.findall(exp, content)
    if mode == 1:
        return match if match else []
    return match[0] if match else None

def java_binder(cpu_usage, mainstack, process):
    detail = detect_keyword(r'\s+at\s+android\.os\.BinderProxy\.transact\(Native Method\)\n\s+(at\s+.+)\n', mainstack)
    if detail:
        return process, detail, None
    return None, None, None

def system_busy(cpu_usage, mainstack, process):
    total = detect_keyword(r'(\d+(?:\.\d+)?)\%\s+TOTAL', cpu_usage)
    iowait = detect_keyword(r'(\d+(?:\.\d+)?)\%\s+iowait', cpu_usage)
    if iowait and float(iowait) >= IOWAIT_THRESHOLD:
        return "iowait", "CPU usage: iowait >= 35%", None
    if total and float(total) >= TOTAL_THRESHOLD:
        topprocess = detect_keyword(r'\s+\d+(?:\.\d+)?\%\s+(?:\d+)?/([\w-]+(?:\.[\w-]+)*)', cpu_usage)
        return topprocess, ("CPU usage: TOTAL >= 40%, " + topprocess + " occupied the most."), None
    return None, None, None

def native_binder(cpu_usage, mainstack, process):
    details = detect_keyword(r'\s+((?:/.+)+(?:\s+\(.+\))?)', mainstack, 1)
    binder = False
    for detail in details:
        if not binder and 'libbinder.so' in detail:
            binder = True
        if binder and 'libbinder.so' not in detail:
            for key in ['libmedia\.so', 'libtvmanager\.so', 'libaudiomanager\.so', 'libgui\.so', 'libsystem_server\.so', 'libchannelmanager\.so']:
                if detect_keyword((r'%s\s+\(.+\)' %(key)), detail):
                    return key.replace("\\", ""), detail, None
    return None, None, None

def native_lock(cpu_usage, mainstack, process):
    details = detect_keyword(r'\s+((?:/.+)+(?:\s+\(.+\))?)', mainstack, 1)
    lock = False
    for detail in details:
        if not lock and 'libc.so (__futex_syscall' in detail:
            lock = True
        if lock and 'libc.so' not in detail:
            return process, detail, None
    return None, None, None

def java_lock(cpu_usage, mainstack, process):
    for exp in [r'\s+-\s+(waiting\s+on\s+<.+>\s+(?:\(a\s+\w+(?:[.$]\w+)+\))?)', r'\s+-\s+(waiting\s+to\s+lock\s+<.+>\s+(?:\(a\s+\w+(?:[.$]\w+)+\))?)']:
        detail = detect_keyword(exp, mainstack)
        if detail:
            return process, re.sub(r'<\w+>', 'addr.', detail), None
    return None, None, None

def mthread_busy(cpu_usage, mainstack, process):
    for exp in [r'\s+(at\s+java\.io\..+)', r'\s+(at\s+android\.database\.sqlite\..+)']:
        detail = detect_keyword(exp, mainstack)
        if detail:
            return process, detail, None
    return None, None, None

def auto_recover(cpu_usage, mainstack, process):
    details = detect_keyword(r'\s+(at\s+\w+(?:[.$]\w+)+(?:\(.+\))?)', mainstack, 1)
    if details and "android.os.MessageQueue.nativePollOnce" in details[0]:
        return None, None, "RECOVER"
    return None, None, None

def detect_basic_info(logcat):
    process, cpu_usage, mainstack = None, None, None
    for minor_part in logcat.split('\n\n'):
        if not process:
            process = detect_string(minor_part, r'^Process:\s+(\w+(?:[.$_]\w+)+)')
        if not cpu_usage and detect_string(minor_part, r'(CPU\s+usage\s+from\s+-?\w+\s+to\s+-?\w+)'):
            cpu_usage = minor_part
        if not mainstack and detect_string(minor_part, r'^(\"main\")\s+'):
            mainstack = minor_part
    return process, cpu_usage, mainstack

METHODS = [java_binder, system_busy, native_binder, native_lock, java_lock, mthread_busy, auto_recover]

def app_anr(logcat, headers):
    process, cpu_usage, mainstack = detect_basic_info(logcat)
    if process and cpu_usage and mainstack:
        for method in METHODS:
            proc, detail, tag = method(cpu_usage, mainstack, process)
            if tag == "RECOVER":
                break
            if proc and detail:
                md5 = gen_hashcode({'issue_owner': proc, 'detail': detail})
                return md5, {'issue_owner': proc, 'detail': detail}, None
    return None, None, None
<|endoftext|>


<commit_before>from database.user import list_users
from framework.database import make_db_connection, close_db_connection
from framework.mail import send_mail

db_conn = make_db_connection()

users = list_users(db_conn, {})

for user in users:
    print('email', user.get('email'))
    if user.get('email'):
        send_mail(
            subject='Welcome to Sagefy',
            recipient=user['email'],
            body="""
            Welcome to Sagefy!

            If you are interested in biweekly updates on Sagefy's progress,
            sign up at http://newsletter.sagefy.org/up

            Thank you!
            """
        )

close_db_connection(db_conn)
<commit_msg>Remove bump script<commit_after><|endoftext|>