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

import cv2
import numpy as np
from sys import float_info

import time

def clip(x, up, low):
    if up < low:
        up, low = low, up
    return max(low, min(up, x))

def cvt_line(rho, theta, shape):
    r, t = rho, theta
    h, w = shape[:2]
    h -= 1
    w -= 1
    sin = np.sin
    cos = np.cos
    if abs(cos(t)) > cos(np.radians(45)):
        x1 = r/cos(t)
        x2 = (rho - h * sin(t))/cos(t)
        x1 = clip(x1, 0, w)
        x2 = clip(x2, 0, w)

        if abs(sin(t)) > float_info.epsilon:  # 0
            y1 = (r - x1 * cos(t))/sin(t)
            y2 = (r - x2 * cos(t))/sin(t)
        else:
            y1 = 0
            y2 = h
        y1 = clip(y1, 0, h)
        y2 = clip(y2, 0, h)
    else:
        y1 = r/sin(t)
        y2 = (r - w * cos(t))/sin(t)
        y1 = clip(y1, 0, h)
        y2 = clip(y2, 0, h)

        if abs(cos(t)) > float_info.epsilon:  # 0
            x1 = (rho - y1 * sin(t))/cos(t)
            x2 = (rho - y2 * sin(t))/cos(t)
        else:
            x1 = 0
            x2 = w
        x1 = clip(x1, 0, w)
        x2 = clip(x2, 0, w)

    x1 = int(round(x1))
    x2 = int(round(x2))
    y1 = int(round(y1))
    y2 = int(round(y2))

    return (x1, y1), (x2, y2)

def mid_point(line, shape):
    rho, theta = line
    pt1, pt2 = cvt_line(rho, theta, shape)

    x1, y1 = pt1
    x2, y2 = pt2

    return (x1 + x2)/2, (y1 + y2)/2

def line_dist(l1, l2, shape):
    m1 = mid_point(l1, shape)
    m2 = mid_point(l2, shape)
    return np.sqrt(
            np.power(m1[0] - m2[0], 2) +
            np.power(m1[1] - m2[1], 2)
            )

def intersection(horizontal, vertical, shape):
    height = shape[0]
    width = shape[1]

    if abs(vertical[1]) > float_info.epsilon:
        vertical1 = 0, vertical[0]/np.sin(vertical[1])
        vertical2 = width, -width/np.tan(vertical[1]) + vertical1[1]
    else:
        vertical1 = vertical[0]/np.cos(vertical[1]), 0
        vertical2 = vertical1[0] - height * np.tan(vertical[1]), height

    horizontal1 = 0, horizontal[0]/np.sin(horizontal[1])
    horizontal2 = width, -width/np.tan(horizontal[1]) + horizontal1[1]

    verticalA = vertical2[1] - vertical1[1]
    verticalB = vertical1[0] - vertical2[0]
    verticalC = verticalA * vertical1[0] + verticalB * vertical1[1]

    horizontalA = horizontal2[1] - horizontal1[1]
    horizontalB = horizontal1[0] - horizontal2[0]
    horizontalC = horizontalA * horizontal1[0] + horizontalB * horizontal1[1]

    det = verticalA * horizontalB - verticalB * horizontalA
    pt = (
            (horizontalB * verticalC - verticalB * horizontalC)/det,
            (verticalA * horizontalC - horizontalA * verticalC)/det,
            )

    return pt

def intersections(topEdge, bottomEdge, leftEdge, rightEdge, shape):
    return (
            intersection(topEdge, leftEdge, shape),
            intersection(topEdge, rightEdge, shape),
            intersection(bottomEdge, rightEdge, shape),
            intersection(bottomEdge, leftEdge, shape),
            )

def unwarp(
        img,
        topEdge,
        bottomEdge,
        leftEdge,
        rightEdge,
        group=False,
        maximize=True,
        dest_size=None,
        ):

    shape = img.shape

    if group:
        start = time.time()
        if maximize:
            ptTopLeft = (np.inf, np.inf)
            ptTopRight  = (-np.inf, np.inf)
            ptBottomLeft = (np.inf, -np.inf)
            ptBottomRight = (-np.inf, -np.inf)
        else:
            ptTopLeft = (-np.inf, -np.inf)
            ptTopRight  = (np.inf, -np.inf)
            ptBottomLeft = (-np.inf, np.inf)
            ptBottomRight = (np.inf, np.inf)
        for t in topEdge:
            for b in bottomEdge:
                for l in leftEdge:
                    for r in rightEdge:
                        tl, tr, br, bl = intersections(
                                t,
                                b,
                                l,
                                r,
                                shape)
                        if maximize:
                            ptTopLeft = (
                                    min(ptTopLeft[0], tl[0]),
                                    min(ptTopLeft[1], tl[1]),
                                    )
                            ptTopRight = (
                                    max(ptTopRight[0], tr[0]),
                                    min(ptTopRight[1], tr[1]),
                                    )
                            ptBottomLeft = (
                                    min(ptBottomLeft[0], bl[0]),
                                    max(ptBottomLeft[1], bl[1]),
                                    )
                            ptBottomRight = (
                                    max(ptBottomRight[0], br[0]),
                                    max(ptBottomRight[1], br[1]),
                                    )
                        else:
                            ptTopLeft = (
                                    max(ptTopLeft[0], tl[0]),
                                    max(ptTopLeft[1], tl[1]),
                                    )
                            ptTopRight = (
                                    min(ptTopRight[0], tr[0]),
                                    max(ptTopRight[1], tr[1]),
                                    )
                            ptBottomLeft = (
                                    max(ptBottomLeft[0], bl[0]),
                                    min(ptBottomLeft[1], bl[1]),
                                    )
                            ptBottomRight = (
                                    min(ptBottomRight[0], br[0]),
                                    min(ptBottomRight[1], br[1]),
                                    )

        end = time.time() - start
        print(f'took {end} secods to find intersections')
    else:
        ptTopLeft, ptTopRight, ptBottomRight, ptBottomLeft = intersections(
                topEdge,
                bottomEdge,
                leftEdge,
                rightEdge,
                shape)

    if dest_size is None:
        def dist(p1, p2):
            return np.sqrt(
                np.power(p1[0] - p2[0], 2) +
                np.power(p1[1] - p2[1], 2)
                )

        width = max(
                dist(ptBottomLeft, ptBottomRight),
                dist(ptTopLeft, ptTopRight))
        height = max(
                dist(ptTopLeft, ptBottomLeft),
                dist(ptTopRight, ptBottomRight))
        width = int(width)
        height = int(height)
    else:
        width, height = dest_size

    src = [ptTopLeft, ptTopRight, ptBottomRight, ptBottomLeft]
    dst = [
            (0, 0),
            (width - 1, 0),
            (width - 1, height - 1),
            (0, height - 1),
            ]

    src = np.asarray(src).astype('float32')
    dst = np.asarray(dst).astype('float32')

    M = cv2.getPerspectiveTransform(src, dst)
    return cv2.warpPerspective(img, M, (width, height))

def find_lines(img, acc_threshold=0.25):
    if len(img.shape) == 3 and img.shape[2] == 3:  # if it's color
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = cv2.GaussianBlur(img, (11, 11), 0)
    img = cv2.adaptiveThreshold(
            img,
            255,
            cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY,
            5,
            2)

    img = cv2.bitwise_not(img)

    # thresh = 127
    # edges = cv2.threshold(img, thresh, 255, cv2.THRESH_BINARY)[1]
    # edges = cv2.Canny(blur, 500, 500, apertureSize=3)

    should_erode = True
    if should_erode:
        element = cv2.getStructuringElement(cv2.MORPH_RECT, (4, 4))
        img = cv2.erode(img, element)

    theta = np.pi/2000
    angle_threshold = 2
    horizontal = cv2.HoughLines(
            img,
            1,
            theta,
            int(acc_threshold * img.shape[1]),
            min_theta=np.radians(90 - angle_threshold),
            max_theta=np.radians(90 + angle_threshold))
    vertical = cv2.HoughLines(
            img,
            1,
            theta,
            int(acc_threshold * img.shape[0]),
            min_theta=np.radians(-angle_threshold),
            max_theta=np.radians(angle_threshold),
            )

    horizontal = list(horizontal) if horizontal is not None else []
    vertical = list(vertical) if vertical is not None else []

    horizontal = [line[0] for line in horizontal]
    vertical = [line[0] for line in vertical]

    horizontal = np.asarray(horizontal)
    vertical = np.asarray(vertical)

    return horizontal, vertical

def group_lines(lines, shape, direction, merge_distance=40):

    def dist(l1, l2):  # auxiliary function
        if direction == 'horizontal':
            return abs(mid_point(l1, shape)[1] - mid_point(l2, shape)[1])
        if direction == 'vertical':
            return abs(mid_point(l1, shape)[0] - mid_point(l2, shape)[0])
        if direction == 'both':
            return line_dist(l1, l2, shape)

    bins = []
    for line in lines:
        m = mid_point(line, shape)
        nearest_bin = None
        for bin in bins:
            nearest_line_dist = min(dist(line, l) for l in bin)
            if nearest_line_dist < merge_distance:
                nearest_bin = bin
                break
        else:
            nearest_bin = []
            bins.append(nearest_bin)
        nearest_bin.append(line)

    bins = sorted(
            bins,
            key=lambda bin: np.average([rho for rho, theta in bin])
            )

    return bins

def merge_lines(bins):
    lines = []
    for bin in bins:
        lines.append((
            np.median([rho for rho, _ in bin]),
            np.median([theta for _, theta in bin]),
            ))

    return lines
<commit_msg>Optimize O(n^4) to O(n^2)<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import numpy as np
from sys import float_info

import time

def clip(x, up, low):
    if up < low:
        up, low = low, up
    return max(low, min(up, x))

def cvt_line(rho, theta, shape):
    r, t = rho, theta
    h, w = shape[:2]
    h -= 1
    w -= 1
    sin = np.sin
    cos = np.cos
    if abs(cos(t)) > cos(np.radians(45)):
        x1 = r/cos(t)
        x2 = (rho - h * sin(t))/cos(t)
        x1 = clip(x1, 0, w)
        x2 = clip(x2, 0, w)

        if abs(sin(t)) > float_info.epsilon:  # 0
            y1 = (r - x1 * cos(t))/sin(t)
            y2 = (r - x2 * cos(t))/sin(t)
        else:
            y1 = 0
            y2 = h
        y1 = clip(y1, 0, h)
        y2 = clip(y2, 0, h)
    else:
        y1 = r/sin(t)
        y2 = (r - w * cos(t))/sin(t)
        y1 = clip(y1, 0, h)
        y2 = clip(y2, 0, h)

        if abs(cos(t)) > float_info.epsilon:  # 0
            x1 = (rho - y1 * sin(t))/cos(t)
            x2 = (rho - y2 * sin(t))/cos(t)
        else:
            x1 = 0
            x2 = w
        x1 = clip(x1, 0, w)
        x2 = clip(x2, 0, w)

    x1 = int(round(x1))
    x2 = int(round(x2))
    y1 = int(round(y1))
    y2 = int(round(y2))

    return (x1, y1), (x2, y2)

def mid_point(line, shape):
    rho, theta = line
    pt1, pt2 = cvt_line(rho, theta, shape)

    x1, y1 = pt1
    x2, y2 = pt2

    return (x1 + x2)/2, (y1 + y2)/2

def line_dist(l1, l2, shape):
    m1 = mid_point(l1, shape)
    m2 = mid_point(l2, shape)
    return np.sqrt(
            np.power(m1[0] - m2[0], 2) +
            np.power(m1[1] - m2[1], 2)
            )

def intersection(horizontal, vertical, shape):
    height = shape[0]
    width = shape[1]

    if abs(vertical[1]) > float_info.epsilon:
        vertical1 = 0, vertical[0]/np.sin(vertical[1])
        vertical2 = width, -width/np.tan(vertical[1]) + vertical1[1]
    else:
        vertical1 = vertical[0]/np.cos(vertical[1]), 0
        vertical2 = vertical1[0] - height * np.tan(vertical[1]), height

    horizontal1 = 0, horizontal[0]/np.sin(horizontal[1])
    horizontal2 = width, -width/np.tan(horizontal[1]) + horizontal1[1]

    verticalA = vertical2[1] - vertical1[1]
    verticalB = vertical1[0] - vertical2[0]
    verticalC = verticalA * vertical1[0] + verticalB * vertical1[1]

    horizontalA = horizontal2[1] - horizontal1[1]
    horizontalB = horizontal1[0] - horizontal2[0]
    horizontalC = horizontalA * horizontal1[0] + horizontalB * horizontal1[1]

    det = verticalA * horizontalB - verticalB * horizontalA
    pt = (
            (horizontalB * verticalC - verticalB * horizontalC)/det,
            (verticalA * horizontalC - horizontalA * verticalC)/det,
            )

    return pt

def intersections(
        topEdge,
        bottomEdge,
        leftEdge,
        rightEdge,
        shape,
        group=False,
        maximize=True,
        ):
    if group:
        if maximize:
            ptTopLeft = (np.inf, np.inf)
            ptTopRight = (-np.inf, np.inf)
            ptBottomLeft = (np.inf, -np.inf)
            ptBottomRight = (-np.inf, -np.inf)
            _max, _min = max, min
        else:
            ptTopLeft = (-np.inf, -np.inf)
            ptTopRight = (np.inf, -np.inf)
            ptBottomLeft = (-np.inf, np.inf)
            ptBottomRight = (np.inf, np.inf)
            _max, _min = min, max

        for t in topEdge:
            for l in leftEdge:
                tl = intersection(t, l, shape)
                ptTopLeft = (
                        _min(ptTopLeft[0], tl[0]),
                        _min(ptTopLeft[1], tl[1]),
                        )
            for r in rightEdge:
                tr = intersection(t, r, shape)
                ptTopRight = (
                        _max(ptTopRight[0], tr[0]),
                        _min(ptTopRight[1], tr[1]),
                        )
        for b in bottomEdge:
            for l in leftEdge:
                bl = intersection(b, l, shape)
                ptBottomLeft = (
                        _min(ptBottomLeft[0], bl[0]),
                        _max(ptBottomLeft[1], bl[1]),
                        )
            for r in rightEdge:
                br = intersection(b, r, shape)
                ptBottomRight = (
                        _max(ptBottomRight[0], br[0]),
                        _max(ptBottomRight[1], br[1]),
                        )
        return (ptTopLeft, ptTopRight, ptBottomRight, ptBottomLeft)
    else:
        return (
                intersection(topEdge, leftEdge, shape),
                intersection(topEdge, rightEdge, shape),
                intersection(bottomEdge, rightEdge, shape),
                intersection(bottomEdge, leftEdge, shape),
                )

def unwarp(
        img,
        topEdge,
        bottomEdge,
        leftEdge,
        rightEdge,
        group=False,
        maximize=True,
        dest_size=None,
        ):

    shape = img.shape

    ptTopLeft, ptTopRight, ptBottomRight, ptBottomLeft = intersections(
            topEdge,
            bottomEdge,
            leftEdge,
            rightEdge,
            shape,
            group,
            maximize)

    if dest_size is None:
        def dist(p1, p2):
            return np.sqrt(
                np.power(p1[0] - p2[0], 2) +
                np.power(p1[1] - p2[1], 2)
                )

        width = max(
                dist(ptBottomLeft, ptBottomRight),
                dist(ptTopLeft, ptTopRight))
        height = max(
                dist(ptTopLeft, ptBottomLeft),
                dist(ptTopRight, ptBottomRight))
        width = int(width)
        height = int(height)
    else:
        width, height = dest_size

    src = [ptTopLeft, ptTopRight, ptBottomRight, ptBottomLeft]
    dst = [
            (0, 0),
            (width - 1, 0),
            (width - 1, height - 1),
            (0, height - 1),
            ]

    src = np.asarray(src).astype('float32')
    dst = np.asarray(dst).astype('float32')

    M = cv2.getPerspectiveTransform(src, dst)
    return cv2.warpPerspective(img, M, (width, height))

def find_lines(img, acc_threshold=0.25):
    if len(img.shape) == 3 and img.shape[2] == 3:  # if it's color
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    img = cv2.GaussianBlur(img, (11, 11), 0)
    img = cv2.adaptiveThreshold(
            img,
            255,
            cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY,
            5,
            2)

    img = cv2.bitwise_not(img)

    # thresh = 127
    # edges = cv2.threshold(img, thresh, 255, cv2.THRESH_BINARY)[1]
    # edges = cv2.Canny(blur, 500, 500, apertureSize=3)

    should_erode = True
    if should_erode:
        element = cv2.getStructuringElement(cv2.MORPH_RECT, (4, 4))
        img = cv2.erode(img, element)

    theta = np.pi/2000
    angle_threshold = 2
    horizontal = cv2.HoughLines(
            img,
            1,
            theta,
            int(acc_threshold * img.shape[1]),
            min_theta=np.radians(90 - angle_threshold),
            max_theta=np.radians(90 + angle_threshold))
    vertical = cv2.HoughLines(
            img,
            1,
            theta,
            int(acc_threshold * img.shape[0]),
            min_theta=np.radians(-angle_threshold),
            max_theta=np.radians(angle_threshold),
            )

    horizontal = list(horizontal) if horizontal is not None else []
    vertical = list(vertical) if vertical is not None else []

    horizontal = [line[0] for line in horizontal]
    vertical = [line[0] for line in vertical]

    horizontal = np.asarray(horizontal)
    vertical = np.asarray(vertical)

    return horizontal, vertical

def group_lines(lines, shape, direction, merge_distance=40):

    def dist(l1, l2):  # auxiliary function
        if direction == 'horizontal':
            return abs(mid_point(l1, shape)[1] - mid_point(l2, shape)[1])
        if direction == 'vertical':
            return abs(mid_point(l1, shape)[0] - mid_point(l2, shape)[0])
        if direction == 'both':
            return line_dist(l1, l2, shape)

    bins = []
    for line in lines:
        m = mid_point(line, shape)
        nearest_bin = None
        for bin in bins:
            nearest_line_dist = min(dist(line, l) for l in bin)
            if nearest_line_dist < merge_distance:
                nearest_bin = bin
                break
        else:
            nearest_bin = []
            bins.append(nearest_bin)
        nearest_bin.append(line)

    bins = sorted(
            bins,
            key=lambda bin: np.average([rho for rho, theta in bin])
            )

    return bins

def merge_lines(bins):
    lines = []
    for bin in bins:
        lines.append((
            np.median([rho for rho, _ in bin]),
            np.median([theta for _, theta in bin]),
            ))

    return lines
<|endoftext|>


<commit_before>from __future__ import unicode_literals, division, absolute_import
import logging
import os
import tempfile

from flexget import plugin
from flexget.event import event

log = logging.getLogger('check_subtitles')

class MetainfoSubs(object):
    """
    Set 'subtitles' field for entries, if they are local video files with subs.
    The field is a list of language codes (3-letter ISO-639-3) for each subtitles
    file found on disk and/or subs track found inside video (for MKVs).
    Special "und" code is for unidentified language (i.e. files without language
    code before extension).
    """

    schema = {'type': 'boolean'}

    def on_task_start(self, task, config):
        try:
            import subliminal
        except ImportError as e:
            log.debug('Error importing Subliminal: %s' % e)
            raise plugin.DependencyError('subliminal', 'subliminal',
                'Subliminal module required. ImportError: %s' % e)
        from dogpile.cache.exception import RegionAlreadyConfigured
        try:
            subliminal.region.configure('dogpile.cache.dbm',
                arguments={'filename': os.path.join(tempfile.gettempdir(), 'cachefile.dbm'),
                           'lock_factory': subliminal.cli.MutexLock})
        except RegionAlreadyConfigured:
            pass
        logging.getLogger("subliminal").setLevel(logging.CRITICAL)
        logging.getLogger("enzyme").setLevel(logging.WARNING)

    def on_task_metainfo(self, task, config):
        # check if explicitly disabled (value set to false)
        if config is False:
            return
        for entry in task.entries:
            entry.register_lazy_func(self.get_subtitles, ['subtitles'])

    def get_subtitles(self, entry):
        if entry.get('subtitles', eval_lazy=False) or not ('location' in entry) or \
                ('$RECYCLE.BIN' in entry['location']) or not os.path.exists(entry['location']):
            return
        import subliminal
        try:
            video = subliminal.scan_video(entry['location'])
            lst = [l.alpha3 for l in video.subtitle_languages]
            if lst:
                entry['subtitles'] = lst
                log.trace('Found subtitles %s for %s' % ('/'.join(lst), entry['title']))
        except Exception as e:
            log.debug('Error checking local subtitles for %s: %s' % (entry['title'], e))

@event('plugin.register')
def register_plugin():
    plugin.register(MetainfoSubs, 'check_subtitles', api_ver=2)
<commit_msg>Update subtitles_check.py<commit_after>from __future__ import unicode_literals, division, absolute_import
import logging
import os
import tempfile

from flexget import plugin
from flexget.event import event

log = logging.getLogger('check_subtitles')

class MetainfoSubs(object):
    """
    Set 'subtitles' field for entries, if they are local video files with subs.
    The field is a list of language codes (3-letter ISO-639-3) for each subtitles
    file found on disk and/or subs track found inside video (for MKVs).
    Special "und" code is for unidentified language (i.e. files without language
    code before extension).
    """

    schema = {'type': 'boolean'}

    def on_task_start(self, task, config):
        try:
            import subliminal
            from subliminal.cli import MutexLock
        except ImportError as e:
            log.debug('Error importing Subliminal: %s' % e)
            raise plugin.DependencyError('subliminal', 'subliminal',
                'Subliminal module required. ImportError: %s' % e)
        from dogpile.cache.exception import RegionAlreadyConfigured
        try:
            subliminal.region.configure('dogpile.cache.dbm',
                arguments={'filename': os.path.join(tempfile.gettempdir(), 'cachefile.dbm'),
                           'lock_factory': MutexLock})
        except RegionAlreadyConfigured:
            pass
        logging.getLogger("subliminal").setLevel(logging.CRITICAL)
        logging.getLogger("enzyme").setLevel(logging.WARNING)

    def on_task_metainfo(self, task, config):
        # check if explicitly disabled (value set to false)
        if config is False:
            return
        for entry in task.entries:
            entry.register_lazy_func(self.get_subtitles, ['subtitles'])

    def get_subtitles(self, entry):
        if entry.get('subtitles', eval_lazy=False) or not ('location' in entry) or \
                ('$RECYCLE.BIN' in entry['location']) or not os.path.exists(entry['location']):
            return
        import subliminal
        try:
            video = subliminal.scan_video(entry['location'])
            lst = [l.alpha3 for l in video.subtitle_languages]
            if lst:
                entry['subtitles'] = lst
                log.trace('Found subtitles %s for %s' % ('/'.join(lst), entry['title']))
        except Exception as e:
            log.debug('Error checking local subtitles for %s: %s' % (entry['title'], e))

@event('plugin.register')
def register_plugin():
    plugin.register(MetainfoSubs, 'check_subtitles', api_ver=2)
<|endoftext|>


<commit_before>"""
.. codeauthor:: Tsuyoshi Hombashi <tsuyoshi.hombashi@gmail.com>
"""

import os.path
from textwrap import indent
from typing import Set, Tuple

from path import Path
from pytablereader.interface import AbstractTableReader
from simplesqlite.model import Integer, Model, Text
from simplesqlite.query import And, Attr, Where
from tabledata import TableData

from .._clrm import bright, green, red, yellow
from .._common import DEFAULT_DUP_COL_HANDLER, ResultLogger
from .._const import MAX_VERBOSITY_LEVEL, PROGRAM_NAME, TABLE_NOT_FOUND_MSG_FORMAT
from .._counter import ResultCounter
from ._ipynb_converter import convert_nb
from ._table_creator import TableCreator

class SourceInfo(Model):
    source_id = Integer(not_null=True)
    dir_name = Text()
    base_name = Text(not_null=True)
    format_name = Text(not_null=True)
    dst_table = Text(not_null=True)
    size = Integer()
    mtime = Integer()

    def get_name(self, verbosity_level):
        if verbosity_level == 0 or self.dir_name is None:
            return self.base_name

        return os.path.join(self.dir_name, self.base_name)

class TableConverter:
    def __init__(
        self,
        logger,
        con,
        symbol_replace_value,
        add_pri_key_name,
        convert_configs,
        index_list,
        is_type_inference,
        is_type_hint_header,
        verbosity_level,
        format_name=None,
        encoding=None,
    ):
        self._logger = logger
        self._con = con
        self._symbol_replace_value = symbol_replace_value
        self._convert_configs = convert_configs
        self._index_list = index_list
        self._is_type_inference = is_type_inference
        self._is_type_hint_header = is_type_hint_header
        self._verbosity_level = verbosity_level
        self._format_name = format_name
        self._encoding = encoding

        self._result_counter = ResultCounter()
        self._result_logger = ResultLogger(
            logger, self._con.schema_extractor, self._result_counter, self._verbosity_level
        )
        self._table_creator = TableCreator(
            logger=self._logger,
            dst_con=con,
            add_pri_key_name=add_pri_key_name,
            result_logger=self._result_logger,
            verbosity_level=verbosity_level,
        )

        SourceInfo.attach(con, is_hidden=True)
        SourceInfo.create()

    def _fetch_source_id(self, source_info):
        where_list = [
            Where("base_name", source_info.base_name),
            Where("format_name", source_info.format_name),
        ]

        if source_info.dir_name:
            where_list.append(Where("dir_name", source_info.dir_name))
        if source_info.size is not None:
            where_list.append(Where("size", source_info.size))
        if source_info.mtime is not None:
            where_list.append(Where("mtime", source_info.mtime))

        return self._con.fetch_value(
            select=Attr("source_id"), table_name=SourceInfo.get_table_name(), where=And(where_list)
        )

    def _fetch_next_source_id(self) -> int:
        source_id = self._con.fetch_value(
            select="MAX({})".format("source_id"), table_name=SourceInfo.get_table_name()
        )

        if source_id is None:
            return 1

        return source_id + 1

    def get_return_code(self) -> int:
        return self._result_counter.get_return_code()

    def get_success_count(self) -> int:
        return self._result_counter.success_count

    def normalize_table(self, table_data: TableData, dup_col_handler=None) -> TableData:
        from pathvalidate import replace_symbol, replace_unprintable_char
        from simplesqlite import SQLiteTableDataSanitizer

        if dup_col_handler is None:
            dup_col_handler = DEFAULT_DUP_COL_HANDLER

        normalized_table_data = SQLiteTableDataSanitizer(
            table_data, dup_col_handler=dup_col_handler, is_type_inference=self._is_type_inference
        ).normalize()

        if self._symbol_replace_value is None:
            return normalized_table_data

        return TableData(
            normalized_table_data.table_name,
            [
                replace_symbol(
                    replace_unprintable_char(header),
                    self._symbol_replace_value,
                    is_replace_consecutive_chars=True,
                    is_strip=True,
                )
                for header in normalized_table_data.headers
            ],
            normalized_table_data.rows,
            dp_extractor=normalized_table_data.dp_extractor,
            type_hints=table_data.dp_extractor.column_type_hints,
        )

    def write_completion_message(self) -> None:
        logger = self._logger

        logger.debug("----- {:s} completed -----".format(PROGRAM_NAME))

        log_list = [
            "source={}".format(
                bright(
                    self._con.fetch_value(
                        select="COUNT(DISTINCT({}))".format("source_id"),
                        table_name=SourceInfo.get_table_name(),
                    )
                )
            )
        ]
        if self.get_success_count() > 0:
            log_list.append(green("success={}".format(bright(str(self.get_success_count())))))
        if self._result_counter.fail_count > 0:
            log_list.append(red("fail={}".format(bright(str(self._result_counter.fail_count)))))
        if self._result_counter.skip_count > 0:
            log_list.append(yellow("skip={}".format(bright(str(self._result_counter.skip_count)))))
        if self._result_counter.created_table_count > 0:
            log_list.append(
                "created-table={}".format(bright(str(self._result_counter.created_table_count)))
            )

        logger.info("converted results: {}".format(", ".join(log_list)))
        database_path_msg = "database path: {:s}".format(
            bright(Path(self._con.database_path).relpath())
        )

        if self.get_success_count() > 0:
            output_format, verbosity_level = self.__get_dump_param()
            logger.info(database_path_msg)

            try:
                logger.debug(
                    "----- database schema -----\n{}".format(
                        indent(
                            self._con.schema_extractor.dumps(
                                output_format=output_format, verbosity_level=verbosity_level
                            ),
                            "    ",
                        )
                    )
                )
            except:  # noqa
                # avoid crashes caused by logging
                pass
        else:
            logger.debug(database_path_msg)

    def _convert_nb(self, nb, source_info):
        success_count = self._result_counter.success_count
        created_table_set = convert_nb(
            logger=self._logger,
            source_info=source_info,
            con=self._con,
            result_logger=self._result_logger,
            nb=nb,
        )

        if self._result_counter.success_count == success_count:
            self._logger.warning(TABLE_NOT_FOUND_MSG_FORMAT.format(source_info.base_name))
            return

        return created_table_set

    def _convert_complex_json(self, json_loader: AbstractTableReader, source_info) -> Set[str]:
        from ._dict_converter import DictConverter

        dict_converter = DictConverter(
            self._logger, self._table_creator, source_info=source_info, index_list=self._index_list
        )

        try:
            dict_converter.to_sqlite_table(json_loader.load_dict(), [])
        except AttributeError:
            pass

        return dict_converter.converted_table_name_set

    def __get_dump_param(self) -> Tuple[str, int]:
        found_ptw = True
        try:
            import pytablewriter  # noqa: W0611
        except ImportError:
            found_ptw = False

        if found_ptw:
            return ("rst_simple_table", self._verbosity_level)

        if self._verbosity_level >= 1:
            return ("text", MAX_VERBOSITY_LEVEL)

        if self._verbosity_level == 0:
            return ("text", 1)

        raise ValueError("invalid verbosity_level: {}".format(self._verbosity_level))
<commit_msg>Remove an unused method<commit_after>"""
.. codeauthor:: Tsuyoshi Hombashi <tsuyoshi.hombashi@gmail.com>
"""

import os.path
from textwrap import indent
from typing import Set, Tuple

from path import Path
from pytablereader.interface import AbstractTableReader
from simplesqlite.model import Integer, Model, Text
from tabledata import TableData

from .._clrm import bright, green, red, yellow
from .._common import DEFAULT_DUP_COL_HANDLER, ResultLogger
from .._const import MAX_VERBOSITY_LEVEL, PROGRAM_NAME, TABLE_NOT_FOUND_MSG_FORMAT
from .._counter import ResultCounter
from ._ipynb_converter import convert_nb
from ._table_creator import TableCreator

class SourceInfo(Model):
    source_id = Integer(not_null=True)
    dir_name = Text()
    base_name = Text(not_null=True)
    format_name = Text(not_null=True)
    dst_table = Text(not_null=True)
    size = Integer()
    mtime = Integer()

    def get_name(self, verbosity_level):
        if verbosity_level == 0 or self.dir_name is None:
            return self.base_name

        return os.path.join(self.dir_name, self.base_name)

class TableConverter:
    def __init__(
        self,
        logger,
        con,
        symbol_replace_value,
        add_pri_key_name,
        convert_configs,
        index_list,
        is_type_inference,
        is_type_hint_header,
        verbosity_level,
        format_name=None,
        encoding=None,
    ):
        self._logger = logger
        self._con = con
        self._symbol_replace_value = symbol_replace_value
        self._convert_configs = convert_configs
        self._index_list = index_list
        self._is_type_inference = is_type_inference
        self._is_type_hint_header = is_type_hint_header
        self._verbosity_level = verbosity_level
        self._format_name = format_name
        self._encoding = encoding

        self._result_counter = ResultCounter()
        self._result_logger = ResultLogger(
            logger, self._con.schema_extractor, self._result_counter, self._verbosity_level
        )
        self._table_creator = TableCreator(
            logger=self._logger,
            dst_con=con,
            add_pri_key_name=add_pri_key_name,
            result_logger=self._result_logger,
            verbosity_level=verbosity_level,
        )

        SourceInfo.attach(con, is_hidden=True)
        SourceInfo.create()

    def _fetch_next_source_id(self) -> int:
        source_id = self._con.fetch_value(
            select="MAX({})".format("source_id"), table_name=SourceInfo.get_table_name()
        )

        if source_id is None:
            return 1

        return source_id + 1

    def get_return_code(self) -> int:
        return self._result_counter.get_return_code()

    def get_success_count(self) -> int:
        return self._result_counter.success_count

    def normalize_table(self, table_data: TableData, dup_col_handler=None) -> TableData:
        from pathvalidate import replace_symbol, replace_unprintable_char
        from simplesqlite import SQLiteTableDataSanitizer

        if dup_col_handler is None:
            dup_col_handler = DEFAULT_DUP_COL_HANDLER

        normalized_table_data = SQLiteTableDataSanitizer(
            table_data, dup_col_handler=dup_col_handler, is_type_inference=self._is_type_inference
        ).normalize()

        if self._symbol_replace_value is None:
            return normalized_table_data

        return TableData(
            normalized_table_data.table_name,
            [
                replace_symbol(
                    replace_unprintable_char(header),
                    self._symbol_replace_value,
                    is_replace_consecutive_chars=True,
                    is_strip=True,
                )
                for header in normalized_table_data.headers
            ],
            normalized_table_data.rows,
            dp_extractor=normalized_table_data.dp_extractor,
            type_hints=table_data.dp_extractor.column_type_hints,
        )

    def write_completion_message(self) -> None:
        logger = self._logger

        logger.debug("----- {:s} completed -----".format(PROGRAM_NAME))

        log_list = [
            "source={}".format(
                bright(
                    self._con.fetch_value(
                        select="COUNT(DISTINCT({}))".format("source_id"),
                        table_name=SourceInfo.get_table_name(),
                    )
                )
            )
        ]
        if self.get_success_count() > 0:
            log_list.append(green("success={}".format(bright(str(self.get_success_count())))))
        if self._result_counter.fail_count > 0:
            log_list.append(red("fail={}".format(bright(str(self._result_counter.fail_count)))))
        if self._result_counter.skip_count > 0:
            log_list.append(yellow("skip={}".format(bright(str(self._result_counter.skip_count)))))
        if self._result_counter.created_table_count > 0:
            log_list.append(
                "created-table={}".format(bright(str(self._result_counter.created_table_count)))
            )

        logger.info("converted results: {}".format(", ".join(log_list)))
        database_path_msg = "database path: {:s}".format(
            bright(Path(self._con.database_path).relpath())
        )

        if self.get_success_count() > 0:
            output_format, verbosity_level = self.__get_dump_param()
            logger.info(database_path_msg)

            try:
                logger.debug(
                    "----- database schema -----\n{}".format(
                        indent(
                            self._con.schema_extractor.dumps(
                                output_format=output_format, verbosity_level=verbosity_level
                            ),
                            "    ",
                        )
                    )
                )
            except:  # noqa
                # avoid crashes caused by logging
                pass
        else:
            logger.debug(database_path_msg)

    def _convert_nb(self, nb, source_info):
        success_count = self._result_counter.success_count
        created_table_set = convert_nb(
            logger=self._logger,
            source_info=source_info,
            con=self._con,
            result_logger=self._result_logger,
            nb=nb,
        )

        if self._result_counter.success_count == success_count:
            self._logger.warning(TABLE_NOT_FOUND_MSG_FORMAT.format(source_info.base_name))
            return

        return created_table_set

    def _convert_complex_json(self, json_loader: AbstractTableReader, source_info) -> Set[str]:
        from ._dict_converter import DictConverter

        dict_converter = DictConverter(
            self._logger, self._table_creator, source_info=source_info, index_list=self._index_list
        )

        try:
            dict_converter.to_sqlite_table(json_loader.load_dict(), [])
        except AttributeError:
            pass

        return dict_converter.converted_table_name_set

    def __get_dump_param(self) -> Tuple[str, int]:
        found_ptw = True
        try:
            import pytablewriter  # noqa: W0611
        except ImportError:
            found_ptw = False

        if found_ptw:
            return ("rst_simple_table", self._verbosity_level)

        if self._verbosity_level >= 1:
            return ("text", MAX_VERBOSITY_LEVEL)

        if self._verbosity_level == 0:
            return ("text", 1)

        raise ValueError("invalid verbosity_level: {}".format(self._verbosity_level))
<|endoftext|>


<commit_before>#!/usr/local/bin/python3
#
# surlybot.py
#
# This is a simple script that will sing the famous "99 bottles of beer" song.  It will default to starting with
# 99 bottles of beer, but you can add optional arguments for changing the beverage type and starting number.
#
# 08/14/2015 - Dan West - Initial Version
# 08/21/2015 - Dan West - Added optional arguments for specifying the beverage type and number to start with
# 09/02/2015 - Dan West - Added a list and loop for the lyrics
#
# Notes
# I just needed something to test Github
#
# Known Issues/ToDo
# - Finish Writing the Script
# - Get drunker

# Import some modules
import sys, time, argparse

# Check for arguments
def arg_check():
    parser = argparse.ArgumentParser(prog='surlybot.py')
    parser.add_argument('-b', '--beverage', default='beer', required=False, action='store', dest='bev_type', help='Requires a valid beverage')
    parser.add_argument('-n', '--number', default='10', required=False, action='store', dest='bottle_num', help='Requires a number greater than 0')
    args = vars(parser.parse_args())
    bev_type = str(args['bev_type'])
    bottle_num = int(args['bottle_num'])
    return bev_type, bottle_num

# Sing the song
def singing(bev_type, bottle_num):
	bev_type = bev_type
	bottle_num = bottle_num
	while bottle_num > 0:

		# Add lyrics to each verse
		verse = []
		verse.append(str(bottle_num) + " bottles of " + str(bev_type) + " on the wall\n")
		verse.append(str(bottle_num) + " bottles of " + str(bev_type) + "\n")
		verse.append("Take one down\n")
		verse.append("Pass it around\n")

		# Take a bottle off the wall
		bottle_num -= 1

		# Now there is one less bottle on the wall
		verse.append(str(bottle_num) + " bottles of " + str(bev_type) + " on the wall\n")

		# Actually sing now
		for line in verse:
			print(line)
			time.sleep(2)
		time.sleep(3)

# Main Function
def main():
	# Check optional arguments
	bev_type, bottle_num = arg_check()

	# Start singing the song
	singing(bev_type, bottle_num)

# Run the main script
main()
<commit_msg>Changd the sleep time at the end of the verse<commit_after>#!/usr/local/bin/python3
#
# surlybot.py
#
# This is a simple script that will sing the famous "99 bottles of beer" song.  It will default to starting with
# 99 bottles of beer, but you can add optional arguments for changing the beverage type and starting number.
#
# 08/14/2015 - Dan West - Initial Version
# 08/21/2015 - Dan West - Added optional arguments for specifying the beverage type and number to start with
# 09/02/2015 - Dan West - Added a list and loop for the lyrics
#
# Notes
# I just needed something to test Github
#
# Known Issues/ToDo
# - Finish Writing the Script
# - Get drunker

# Import some modules
import sys, time, argparse

# Check for arguments
def arg_check():
    parser = argparse.ArgumentParser(prog='surlybot.py')
    parser.add_argument('-b', '--beverage', default='beer', required=False, action='store', dest='bev_type', help='Requires a valid beverage')
    parser.add_argument('-n', '--number', default='10', required=False, action='store', dest='bottle_num', help='Requires a number greater than 0')
    args = vars(parser.parse_args())
    bev_type = str(args['bev_type'])
    bottle_num = int(args['bottle_num'])
    return bev_type, bottle_num

# Sing the song
def singing(bev_type, bottle_num):
	bev_type = bev_type
	bottle_num = bottle_num
	while bottle_num > 0:

		# Add lyrics to each verse
		verse = []
		verse.append(str(bottle_num) + " bottles of " + str(bev_type) + " on the wall\n")
		verse.append(str(bottle_num) + " bottles of " + str(bev_type) + "\n")
		verse.append("Take one down\n")
		verse.append("Pass it around\n")

		# Take a bottle off the wall
		bottle_num -= 1

		# Now there is one less bottle on the wall
		verse.append(str(bottle_num) + " bottles of " + str(bev_type) + " on the wall\n")

		# Actually sing now
		for line in verse:
			print(line)
			time.sleep(2)
		time.sleep(1)

# Main Function
def main():
	# Check optional arguments
	bev_type, bottle_num = arg_check()

	# Start singing the song
	singing(bev_type, bottle_num)

# Run the main script
main()
<|endoftext|>


<commit_before>import os
import sys
import configparser
MS_WD = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# Append .. to sys path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Add the storage dir to the sys.path. Allows import of sql_driver module
if os.path.join(MS_WD, 'storage') not in sys.path:
    sys.path.insert(0, os.path.join(MS_WD, 'storage'))
# Add the libs dir to the sys.path. Allows import of common module
if os.path.join(MS_WD, 'libs') not in sys.path:
    sys.path.insert(0, os.path.join(MS_WD, 'libs'))
import multiscanner
import common
import sql_driver as database

from celery import Celery

DEFAULTCONF = {
    'protocol': 'pyamqp',
    'host': 'localhost',
    'user': 'guest',
    'password': '',
    'vhost': '/',
}

config_object = configparser.SafeConfigParser()
config_object.optionxform = str
configfile = common.get_api_config_path(multiscanner.CONFIG)
config_object.read(configfile)
config = common.parse_config(config_object)
worker_config = config.get('celery')
db_config = config.get('Database')

app = Celery(broker='{0}://{1}:{2}@{3}/{4}'.format(
    worker_config.get('protocol'),
    worker_config.get('user'),
    worker_config.get('password'),
    worker_config.get('host'),
    worker_config.get('vhost'),
))
db = database.Database(config=db_config)

@app.task
def multiscanner_celery(file_, original_filename, task_id, file_hash, config=multiscanner.CONFIG):
    '''
    TODO: Figure out how to do batching.
    This function essentially takes in a file list and runs
    multiscanner on them. Results are stored in the
    storage configured in storage.ini.

    Usage:
    from celery_worker import multiscanner_celery
    multiscanner_celery.delay([list, of, files, to, scan])
    '''
    # Initialize the connection to the task DB
    db.init_db()

    # Get the storage config
    storage_conf = multiscanner.common.get_storage_config_path(config)
    storage_handler = multiscanner.storage.StorageHandler(configfile=storage_conf)

    resultlist = multiscanner.multiscan([file_], configfile=config)
    results = multiscanner.parse_reports(resultlist, python=True)
    # Use the original filename as the value for the filename
    # in the report (instead of the tmp path assigned to the file
    # by the REST API)
    results[original_filename] = results[file_]
    del results[file_]

    # Save the report to storage
    storage_handler.store(results, wait=False)
    storage_handler.close()

    # Update the task DB to reflect that the task is done
    db.update_task(
        task_id=task_id,
        task_status='Complete',
        report_id=file_hash
    )

    return results

if __name__ == '__main__':
    app.start()
<commit_msg>Add run instructions to celery worker<commit_after>'''
This is the multiscanner celery worker. To initialize a worker node run:
$ celery -A celery_worker worker
from the utils/ directory.
'''

import os
import sys
import configparser
MS_WD = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# Append .. to sys path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Add the storage dir to the sys.path. Allows import of sql_driver module
if os.path.join(MS_WD, 'storage') not in sys.path:
    sys.path.insert(0, os.path.join(MS_WD, 'storage'))
# Add the libs dir to the sys.path. Allows import of common module
if os.path.join(MS_WD, 'libs') not in sys.path:
    sys.path.insert(0, os.path.join(MS_WD, 'libs'))
import multiscanner
import common
import sql_driver as database

from celery import Celery

DEFAULTCONF = {
    'protocol': 'pyamqp',
    'host': 'localhost',
    'user': 'guest',
    'password': '',
    'vhost': '/',
}

config_object = configparser.SafeConfigParser()
config_object.optionxform = str
configfile = common.get_api_config_path(multiscanner.CONFIG)
config_object.read(configfile)
config = common.parse_config(config_object)
worker_config = config.get('celery')
db_config = config.get('Database')

app = Celery(broker='{0}://{1}:{2}@{3}/{4}'.format(
    worker_config.get('protocol'),
    worker_config.get('user'),
    worker_config.get('password'),
    worker_config.get('host'),
    worker_config.get('vhost'),
))
db = database.Database(config=db_config)

@app.task
def multiscanner_celery(file_, original_filename, task_id, file_hash, config=multiscanner.CONFIG):
    '''
    TODO: Figure out how to do batching.
    This function essentially takes in a file list and runs
    multiscanner on them. Results are stored in the
    storage configured in storage.ini.

    Usage:
    from celery_worker import multiscanner_celery
    multiscanner_celery.delay([list, of, files, to, scan])
    '''
    # Initialize the connection to the task DB
    db.init_db()

    # Get the storage config
    storage_conf = multiscanner.common.get_storage_config_path(config)
    storage_handler = multiscanner.storage.StorageHandler(configfile=storage_conf)

    resultlist = multiscanner.multiscan([file_], configfile=config)
    results = multiscanner.parse_reports(resultlist, python=True)
    # Use the original filename as the value for the filename
    # in the report (instead of the tmp path assigned to the file
    # by the REST API)
    results[original_filename] = results[file_]
    del results[file_]

    # Save the report to storage
    storage_handler.store(results, wait=False)
    storage_handler.close()

    # Update the task DB to reflect that the task is done
    db.update_task(
        task_id=task_id,
        task_status='Complete',
        report_id=file_hash
    )

    return results

if __name__ == '__main__':
    app.start()
<|endoftext|>


<commit_before>from __future__ import print_function

import math
import operator as op
from collections import OrderedDict

def append(*values):
    return reduce(op.add, values)

def dict_put(data_dict, key, value):
    data = data_dict.copy()
    data[key] = value
    return data

def dict_get(data_dict, key):
    return data_dict[key]

def dict_has_key(data_dict, key):
    return data_dict.get(key) is not None

class BotcenterDSLPrimitives(object):

    MATH = vars(math)

    OPERATORS = {
        '+': op.add,
        '-': op.sub,
        '*': op.mul,
        '/': lambda x, y: x / y,
        '>': op.gt,
        '<': op.lt,
        '>=': op.ge,
        '<=': op.le,
        '=': op.eq,
        'abs': abs,
        'equal?': op.eq,
        'not': op.not_,
    }

    LIST_OPERATIONS = {
        'append': append,
        'head': lambda x: x[0],
        'tail': lambda x: x[1:],
        'length': len,
        'list': lambda *x: list(x),
        'map': lambda f, l: list(map(f, l)),
        'max': max,
        'min': min,
    }

    DICT_OPERATIONS = {
        'make-dict': lambda: OrderedDict(),
        'put': dict_put,
        'get': dict_get,
        'has-key': dict_has_key
    }

    SIDE_EFFECTS = {
        'print': print
    }

    @classmethod
    def populate_environment(cls, environment):

        environment.add_primitives(cls.MATH)
        environment.add_primitives(cls.OPERATORS)
        environment.add_primitives(cls.LIST_OPERATIONS)
        environment.add_primitives(cls.DICT_OPERATIONS)
        environment.add_primitives(cls.SIDE_EFFECTS)
        return environment
<commit_msg>Added in-list primitive<commit_after>from __future__ import print_function

import math
import operator as op
from collections import OrderedDict

def append(*values):
    return reduce(op.add, values)

def dict_put(data_dict, key, value):
    data = data_dict.copy()
    data[key] = value
    return data

def dict_get(data_dict, key):
    return data_dict[key]

def dict_has_key(data_dict, key):
    return data_dict.get(key) is not None

class BotcenterDSLPrimitives(object):

    MATH = vars(math)

    OPERATORS = {
        '+': op.add,
        '-': op.sub,
        '*': op.mul,
        '/': lambda x, y: x / y,
        '>': op.gt,
        '<': op.lt,
        '>=': op.ge,
        '<=': op.le,
        '=': op.eq,
        'abs': abs,
        'equal?': op.eq,
        'not': op.not_,
    }

    LIST_OPERATIONS = {
        'append': append,
        'head': lambda x: x[0],
        'tail': lambda x: x[1:],
        'length': len,
        'list': lambda *x: list(x),
        'map': lambda f, l: list(map(f, l)),
        'max': max,
        'min': min,
        'in-list': lambda lst, elem: elem in lst
    }

    DICT_OPERATIONS = {
        'make-dict': lambda: OrderedDict(),
        'put': dict_put,
        'get': dict_get,
        'has-key': dict_has_key
    }

    SIDE_EFFECTS = {
        'print': print
    }

    @classmethod
    def populate_environment(cls, environment):

        environment.add_primitives(cls.MATH)
        environment.add_primitives(cls.OPERATORS)
        environment.add_primitives(cls.LIST_OPERATIONS)
        environment.add_primitives(cls.DICT_OPERATIONS)
        environment.add_primitives(cls.SIDE_EFFECTS)
        return environment
<|endoftext|>


<commit_before><commit_msg>Update GameBoy.py<commit_after><|endoftext|>


<commit_before># Copyright Hybrid Logic Ltd.  See LICENSE file for details.

"""
Tests for ``flocker.provision._digitalocean``.
"""
import copy
import os

try:
    import pyocean
except ImportError:
    PYOCEAN_INSTALLED=False
else:
    PYOCEAN_INSTALLED=True

from twisted.trial.unittest import SynchronousTestCase, SkipTest, skipUnless

from flocker.provision._digitalocean import (
    set_droplet_kernel, retry_if_pending, latest_droplet_kernel,
    kernel_from_digitalocean_version, DIGITALOCEAN_KERNEL)
from flocker.testtools import random_name

TESTING_DROPLET_ATTRIBUTES = {
    'region': 'lon1',
    'size': '8gb',
    'image': 'fedora-20-x64'
}

@skipUnless(PYOCEAN_INSTALLED, "pyocean not installed")
def client_from_environment():
    token = os.environ.get('DIGITALOCEAN_TOKEN')
    if token is None:
        raise SkipTest(
            'A DIGITALOCEAN_TOKEN environment variable is required to run '
            'these tests.')

    return pyocean.DigitalOcean(token)

def droplet_for_test(test_case, client):
    droplet_attributes = copy.deepcopy(TESTING_DROPLET_ATTRIBUTES)
    droplet_attributes['name'] = (
        test_case.id().replace('_', '-')
        + '-'
        + random_name()
    )
    droplet = retry_if_pending(client.droplet.create, droplet_attributes)
    test_case.addCleanup(retry_if_pending, droplet.destroy)
    return droplet

class SetDropletKernelTests(SynchronousTestCase):
    """
    Tests for ``set_droplet_kernel``.

    These tests are designed to interact with live DigitalOcean droplets.

    You must supply a DigitalOcean V2 API token by setting
    ``DIGITALOCEAN_TOKEN`` in the environment before running these tests.
    """
    def setUp(self):
        """
        Set up a test droplet and destroy it after the test.
        """
        self.client = client_from_environment()
        self.droplet = droplet_for_test(self, self.client)

    def test_success(self):
        """
        ``set_droplet_kernel`` assigns the supplied kernel to the droplet.
        to the droplet, returning the selected DigitalOcean kernel instance.
        """
        expected_kernel = DIGITALOCEAN_KERNEL
        set_droplet_kernel(self.droplet, expected_kernel)

        # Need to query again for the droplet after updating its kernel
        updated_droplet = self.client.droplet.get(self.droplet.id)
        # Pyocean wraps kernel attributes in its ``Image`` class...which makes
        # no sense and then uses a ``dict`` for the ``droplet.kernel``
        # attribute, so they can't be directly compared. It would be better if
        # both used a dedicated ``Kernel`` type which could easily be
        # compared.
        # See: https://github.com/flowfree/pyocean/issues/2
        actual_kernel = kernel_from_digitalocean_version(
            updated_droplet.kernel['version']
        )
        self.assertEqual(expected_kernel, actual_kernel)

class LatestDropletKernelTests(SynchronousTestCase):
    """
    Tests for ``latest_droplet_kernel``.

    These tests are designed to interact with live DigitalOcean droplets.

    You must supply a DigitalOcean V2 API token by setting
    ``DIGITALOCEAN_TOKEN`` in the environment before running these tests.
    """
    def setUp(self):
        """
        Set up a test droplet and destroy it after the test.
        """
        self.client = client_from_environment()
        self.droplet = droplet_for_test(self, self.client)

    def test_success(self):
        """
        ``latest_droplet_kernel`` should return the same kernel that we ask
        users to install in our documentation.
        """
        expected_kernel = DIGITALOCEAN_KERNEL
        actual_kernel = latest_droplet_kernel(self.droplet, 'fc20', 'x86_64')

        self.assertEqual(expected_kernel, actual_kernel)
<commit_msg>Wrong guess<commit_after># Copyright Hybrid Logic Ltd.  See LICENSE file for details.

"""
Tests for ``flocker.provision._digitalocean``.
"""
import copy
import os
from unittest import skipUnless

try:
    import pyocean
except ImportError:
    PYOCEAN_INSTALLED=False
else:
    PYOCEAN_INSTALLED=True

from twisted.trial.unittest import SynchronousTestCase, SkipTest

from flocker.provision._digitalocean import (
    set_droplet_kernel, retry_if_pending, latest_droplet_kernel,
    kernel_from_digitalocean_version, DIGITALOCEAN_KERNEL)
from flocker.testtools import random_name

TESTING_DROPLET_ATTRIBUTES = {
    'region': 'lon1',
    'size': '8gb',
    'image': 'fedora-20-x64'
}

@skipUnless(PYOCEAN_INSTALLED, "pyocean not installed")
def client_from_environment():
    token = os.environ.get('DIGITALOCEAN_TOKEN')
    if token is None:
        raise SkipTest(
            'A DIGITALOCEAN_TOKEN environment variable is required to run '
            'these tests.')

    return pyocean.DigitalOcean(token)

def droplet_for_test(test_case, client):
    droplet_attributes = copy.deepcopy(TESTING_DROPLET_ATTRIBUTES)
    droplet_attributes['name'] = (
        test_case.id().replace('_', '-')
        + '-'
        + random_name()
    )
    droplet = retry_if_pending(client.droplet.create, droplet_attributes)
    test_case.addCleanup(retry_if_pending, droplet.destroy)
    return droplet

class SetDropletKernelTests(SynchronousTestCase):
    """
    Tests for ``set_droplet_kernel``.

    These tests are designed to interact with live DigitalOcean droplets.

    You must supply a DigitalOcean V2 API token by setting
    ``DIGITALOCEAN_TOKEN`` in the environment before running these tests.
    """
    def setUp(self):
        """
        Set up a test droplet and destroy it after the test.
        """
        self.client = client_from_environment()
        self.droplet = droplet_for_test(self, self.client)

    def test_success(self):
        """
        ``set_droplet_kernel`` assigns the supplied kernel to the droplet.
        to the droplet, returning the selected DigitalOcean kernel instance.
        """
        expected_kernel = DIGITALOCEAN_KERNEL
        set_droplet_kernel(self.droplet, expected_kernel)

        # Need to query again for the droplet after updating its kernel
        updated_droplet = self.client.droplet.get(self.droplet.id)
        # Pyocean wraps kernel attributes in its ``Image`` class...which makes
        # no sense and then uses a ``dict`` for the ``droplet.kernel``
        # attribute, so they can't be directly compared. It would be better if
        # both used a dedicated ``Kernel`` type which could easily be
        # compared.
        # See: https://github.com/flowfree/pyocean/issues/2
        actual_kernel = kernel_from_digitalocean_version(
            updated_droplet.kernel['version']
        )
        self.assertEqual(expected_kernel, actual_kernel)

class LatestDropletKernelTests(SynchronousTestCase):
    """
    Tests for ``latest_droplet_kernel``.

    These tests are designed to interact with live DigitalOcean droplets.

    You must supply a DigitalOcean V2 API token by setting
    ``DIGITALOCEAN_TOKEN`` in the environment before running these tests.
    """
    def setUp(self):
        """
        Set up a test droplet and destroy it after the test.
        """
        self.client = client_from_environment()
        self.droplet = droplet_for_test(self, self.client)

    def test_success(self):
        """
        ``latest_droplet_kernel`` should return the same kernel that we ask
        users to install in our documentation.
        """
        expected_kernel = DIGITALOCEAN_KERNEL
        actual_kernel = latest_droplet_kernel(self.droplet, 'fc20', 'x86_64')

        self.assertEqual(expected_kernel, actual_kernel)
<|endoftext|>


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

"""
This module provides a session ticket mechanism.

The implemented mechanism is a subset of session tickets as proposed for
TLS in RFC 5077.

The format of a 112-byte ticket is:
 +------------+------------------+--------------+
 | 16-byte IV | 64-byte E(state) | 32-byte HMAC |
 +------------+------------------+--------------+

The 64-byte encrypted state contains:
 +-------------------+--------------------+--------------------+-------------+
 | 4-byte issue date | 18-byte identifier | 32-byte master key | 10-byte pad |
 +-------------------+--------------------+--------------------+-------------+
"""

import os
import time
import const
import pickle
import base64
import struct
import datetime

from Crypto.Cipher import AES
from Crypto.Hash import HMAC
from Crypto.Hash import SHA256

import obfsproxy.common.log as logging

import mycrypto
import util

log = logging.get_obfslogger()

# Length of the IV which is used for AES-CBC.
IV_LENGTH = 16

# Length of the HMAC used to authenticate the ticket.
HMAC_KEY_LENGTH = 32

# Length of the AES key used to encrypt the ticket.
AES_KEY_LENGTH = 16

# Must be a multiple of 16 bytes due to AES' block size.
IDENTIFIER = "ScrambleSuitTicket"

HMACKey = AESKey = creationTime = None

def storeNewTicket( masterKey, ticket, bridge ):
    """
    Store a new session ticket and the according master key for future use.

    The given data is pickled and stored in the global ticket dictionary.  If
    there already is a ticket for the given `bridge', it is overwritten.
    """

    assert len(masterKey) == const.MASTER_KEY_LENGTH
    assert len(ticket) == const.TICKET_LENGTH

    log.debug("Storing newly received ticket in `%s'." % const.TICKET_FILE)

    # Add a new (key, ticket) tuple with the given bridge as hash key.
    tickets = dict()
    content = util.readFromFile(const.TICKET_FILE)
    if (content is not None) and (len(content) > 0):
        tickets = pickle.loads(content)

    tickets[bridge] = (masterKey, ticket)
    util.writeToFile(pickle.dumps(tickets), const.TICKET_FILE)

def findStoredTicket( bridge, fileName=const.TICKET_FILE ):
    """
    Retrieve a previously stored ticket from the ticket dictionary.

    The global ticket dictionary is loaded and the given `bridge' is used to
    look up the ticket and the master key.  If the ticket dictionary does not
    exist (yet) or the ticket data could not be found, `None' is returned.
    """

    assert bridge

    log.debug("Attempting to read master key and ticket from file `%s'." %
              fileName)

    if not os.path.exists(fileName):
        return None

    # Load the ticket hash table from file.
    blurb = util.readFromFile(fileName)
    if (blurb is None) or (len(blurb) == 0):
        return None
    tickets = pickle.loads(blurb)

    try:
        masterKey, ticket = tickets[bridge]
    except KeyError:
        log.info("Found no ticket for bridge `%s'." % str(bridge))
        return None

    return (masterKey, ticket)

def rotateKeys( ):
    """
    Rotate the keys used to encrypt and authenticate session tickets.

    After new keys (an AES and a HMAC key) were created, the old keys are still
    kept for a period of seven days to verify (but not to issue) session
    tickets issues by the old keys.
    """

    global HMACKey
    global AESKey
    global creationTime

    log.info("Attempting to rotate session ticket keys.")

    # Generate and load fresh keys.
    HMACKey = mycrypto.strong_random(HMAC_KEY_LENGTH)
    AESKey = mycrypto.strong_random(AES_KEY_LENGTH)
    creationTime = int(time.time())

    try:
        with open(const.KEY_STORE, 'w') as fd:
            pickle.dump([creationTime, HMACKey, AESKey], fd)
    except IOError as err:
        log.error("Error writing ticket key file to `%s'." % err)

def loadKeys( ):
    """
    Load the keys used to encrypt and authenticate session tickets.

    The keys are loaded from file and stored in global variables so they can be
    accessed from different functions.
    """

    global HMACKey
    global AESKey
    global creationTime

    log.info("Attempting to read ticket keys k_S from file `%s'." %
             const.KEY_STORE)

    # If the key store does not exist (yet), it must be created.
    if not os.path.exists(const.KEY_STORE):
        rotateKeys()
        return

    try:
        with open(const.KEY_STORE, 'r') as fd:
            creationTime, HMACKey, AESKey = pickle.load(fd)
    except IOError as err:
        log.error("Error reading ticket key file from `%s'." % err)

def checkKeys( ):
    """
    Check whether the encryption and authentication keys are defined and valid.

    If the keys are not defined, they are loaded from file by calling
    `loadKeys()'.  If they are expired and no longer valid, the keys are
    rotated by calling `rotateKeys()'.
    """

    if (HMACKey is None) or (AESKey is None):
        loadKeys()

    if (int(time.time()) - creationTime) > const.KEY_ROTATION_TIME:
        rotateKeys()

def decrypt( ticket ):
    """
    Decrypts, verifies and returns the given `ticket'.

    First, the HMAC over the ticket is verified.  If it is valid, the ticket is
    decrypted.  Finally, a `ProtocolState()' object containing the master key
    and the ticket's issue date is returned.  If any of these steps fail,
    `None' is returned.
    """

    assert (ticket is not None) and (len(ticket) == const.TICKET_LENGTH)

    global HMACKey
    global AESKey
    global creationTime

    log.debug("Attempting to decrypt and verify ticket.")

    checkKeys()

    # Verify the ticket's authenticity before decrypting.
    hmac = HMAC.new(HMACKey, ticket[0:80], digestmod=SHA256).digest()
    if hmac != ticket[80:const.TICKET_LENGTH]:
        log.debug("The ticket's HMAC is invalid.  Probably not a ticket.")
        return None

    # Decrypt the ticket to extract the state information.
    aes = AES.new(AESKey, mode=AES.MODE_CBC, IV=ticket[0:IV_LENGTH])
    plainTicket = aes.decrypt(ticket[IV_LENGTH:80])

    issueDate = struct.unpack('I', plainTicket[0:4])[0]
    identifier = plainTicket[4:22]
    masterKey = plainTicket[22:54]

    if not (identifier == IDENTIFIER):
        log.error("The ticket's HMAC is valid but the identifier is invalid.  "
                  "The ticket could be corrupt.")
        return None

    return ProtocolState(masterKey, issueDate=issueDate)

class ProtocolState( object ):

    """
    Defines a ScrambleSuit protocol state contained in a session ticket.

    A protocol state is essentially a master key which can then be used by the
    server to derive session keys.  Besides, a state object contains an issue
    date which specifies the expiry date of a ticket.  This class contains
    methods to check the expiry status of a ticket and to dump it in its raw
    form.
    """

    def __init__( self, masterKey, issueDate=int(time.time()) ):
        """
        The constructor of the `ProtocolState' class.

        The four class variables are initialised.
        """

        self.identifier = IDENTIFIER
        self.masterKey = masterKey
        self.issueDate = issueDate
        # Pad to multiple of 16 bytes to match AES' block size.
        self.pad = "\0\0\0\0\0\0\0\0\0\0"

    def isValid( self ):
        """
        Verifies the expiry date of the object's issue date.

        If the expiry date is not yet reached and the protocol state is still
        valid, `True' is returned.  If the protocol state has expired, `False'
        is returned.
        """

        assert self.issueDate

        lifetime = int(time.time()) - self.issueDate
        if lifetime > const.SESSION_TICKET_LIFETIME:
            log.debug("The ticket is invalid and expired %s ago." %
                      str(datetime.timedelta(seconds=
                      (lifetime - const.SESSION_TICKET_LIFETIME))))
            return False

        log.debug("The ticket is still valid for %s." %
                  str(datetime.timedelta(seconds=
                  (const.SESSION_TICKET_LIFETIME - lifetime))))
        return True

    def __repr__( self ):
        """
        Return a raw string representation of the object's protocol state.

        The length of the returned representation is exactly 64 bytes; a
        multiple of AES' 16-byte block size.  That makes it suitable to be
        encrypted using AES-CBC.
        """

        return struct.pack('I', self.issueDate) + self.identifier + \
                           self.masterKey + self.pad

class SessionTicket( object ):

    """
    Encrypts and authenticates an encapsulated `ProtocolState()' object.

    This class implements a session ticket which can be redeemed by clients.
    The class contains methods to initialise and issue session tickets.
    """

    def __init__( self, masterKey ):
        """
        The constructor of the `SessionTicket()' class.

        The class variables are initialised and the validity of the symmetric
        keys for the session tickets is checked.
        """

        assert (masterKey is not None) and \
               len(masterKey) == const.MASTER_KEY_LENGTH

        checkKeys()

        # Initialisation vector for AES-CBC.
        self.IV = mycrypto.strong_random(IV_LENGTH)

        # The server's (encrypted) protocol state.
        self.state = ProtocolState(masterKey)

        # AES and HMAC keys to encrypt and authenticate the ticket.
        self.symmTicketKey = AESKey
        self.hmacTicketKey = HMACKey

    def issue( self ):
        """
        Returns a ready-to-use session ticket after prior initialisation.

        After the `SessionTicket()' class was initialised with a master key,
        this method encrypts and authenticates the protocol state and returns
        the final result which is ready to be sent over the wire.
        """

        self.state.issueDate = int(time.time())

        # Encrypt the protocol state.
        aes = AES.new(self.symmTicketKey, mode=AES.MODE_CBC, IV=self.IV)
        state = repr(self.state)
        assert (len(state) % AES.block_size) == 0
        cryptedState = aes.encrypt(state)

        # Authenticate the encrypted state and the IV.
        hmac = HMAC.new(self.hmacTicketKey,
                        self.IV + cryptedState, digestmod=SHA256).digest()

        finalTicket = self.IV + cryptedState + hmac
        log.debug("Returning %d-byte ticket." % len(finalTicket))

        return finalTicket

# Alias class name in order to provide a more intuitive API.
new = SessionTicket

# Give ScrambleSuit server operators a way to manually issue new session
# tickets for out-of-band distribution.
if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("ticket_file", type=str, help="The file, the newly "
                        "issued ticket is written to.")
    args = parser.parse_args()

    print "[+] Generating new session ticket."
    masterKey = mycrypto.strong_random(const.MASTER_KEY_LENGTH)
    ticketObj = SessionTicket(masterKey)
    ticket = ticketObj.issue()

    print "[+] Writing new session ticket to `%s'." % args.ticket_file
    util.writeToFile(base64.b32encode(masterKey + ticket) + '\n',
                     args.ticket_file)

    print "[+] Success."
<commit_msg>Adapt the manual ticket issuing mechanism so it works with the improved dictionary-based ticket store.<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
This module provides a session ticket mechanism.

The implemented mechanism is a subset of session tickets as proposed for
TLS in RFC 5077.

The format of a 112-byte ticket is:
 +------------+------------------+--------------+
 | 16-byte IV | 64-byte E(state) | 32-byte HMAC |
 +------------+------------------+--------------+

The 64-byte encrypted state contains:
 +-------------------+--------------------+--------------------+-------------+
 | 4-byte issue date | 18-byte identifier | 32-byte master key | 10-byte pad |
 +-------------------+--------------------+--------------------+-------------+
"""

import os
import time
import const
import pickle
import base64
import struct
import datetime

from Crypto.Cipher import AES
from Crypto.Hash import HMAC
from Crypto.Hash import SHA256
from twisted.internet.address import IPv4Address

import obfsproxy.common.log as logging

import mycrypto
import util

log = logging.get_obfslogger()

# Length of the IV which is used for AES-CBC.
IV_LENGTH = 16

# Length of the HMAC used to authenticate the ticket.
HMAC_KEY_LENGTH = 32

# Length of the AES key used to encrypt the ticket.
AES_KEY_LENGTH = 16

# Must be a multiple of 16 bytes due to AES' block size.
IDENTIFIER = "ScrambleSuitTicket"

HMACKey = AESKey = creationTime = None

def storeNewTicket( masterKey, ticket, bridge ):
    """
    Store a new session ticket and the according master key for future use.

    The given data is pickled and stored in the global ticket dictionary.  If
    there already is a ticket for the given `bridge', it is overwritten.
    """

    assert len(masterKey) == const.MASTER_KEY_LENGTH
    assert len(ticket) == const.TICKET_LENGTH

    log.debug("Storing newly received ticket in `%s'." % const.TICKET_FILE)

    # Add a new (key, ticket) tuple with the given bridge as hash key.
    tickets = dict()
    content = util.readFromFile(const.TICKET_FILE)
    if (content is not None) and (len(content) > 0):
        tickets = pickle.loads(content)

    tickets[bridge] = (masterKey, ticket)
    util.writeToFile(pickle.dumps(tickets), const.TICKET_FILE)

def findStoredTicket( bridge, fileName=const.TICKET_FILE ):
    """
    Retrieve a previously stored ticket from the ticket dictionary.

    The global ticket dictionary is loaded and the given `bridge' is used to
    look up the ticket and the master key.  If the ticket dictionary does not
    exist (yet) or the ticket data could not be found, `None' is returned.
    """

    assert bridge

    log.debug("Attempting to read master key and ticket from file `%s'." %
              fileName)

    if not os.path.exists(fileName):
        return None

    # Load the ticket hash table from file.
    blurb = util.readFromFile(fileName)
    if (blurb is None) or (len(blurb) == 0):
        return None
    tickets = pickle.loads(blurb)

    try:
        masterKey, ticket = tickets[bridge]
    except KeyError:
        log.info("Found no ticket for bridge `%s'." % str(bridge))
        return None

    return (masterKey, ticket)

def rotateKeys( ):
    """
    Rotate the keys used to encrypt and authenticate session tickets.

    After new keys (an AES and a HMAC key) were created, the old keys are still
    kept for a period of seven days to verify (but not to issue) session
    tickets issues by the old keys.
    """

    global HMACKey
    global AESKey
    global creationTime

    log.info("Attempting to rotate session ticket keys.")

    # Generate and load fresh keys.
    HMACKey = mycrypto.strong_random(HMAC_KEY_LENGTH)
    AESKey = mycrypto.strong_random(AES_KEY_LENGTH)
    creationTime = int(time.time())

    try:
        with open(const.KEY_STORE, 'w') as fd:
            pickle.dump([creationTime, HMACKey, AESKey], fd)
    except IOError as err:
        log.error("Error writing ticket key file to `%s'." % err)

def loadKeys( ):
    """
    Load the keys used to encrypt and authenticate session tickets.

    The keys are loaded from file and stored in global variables so they can be
    accessed from different functions.
    """

    global HMACKey
    global AESKey
    global creationTime

    log.info("Attempting to read ticket keys k_S from file `%s'." %
             const.KEY_STORE)

    # If the key store does not exist (yet), it must be created.
    if not os.path.exists(const.KEY_STORE):
        rotateKeys()
        return

    try:
        with open(const.KEY_STORE, 'r') as fd:
            creationTime, HMACKey, AESKey = pickle.load(fd)
    except IOError as err:
        log.error("Error reading ticket key file from `%s'." % err)

def checkKeys( ):
    """
    Check whether the encryption and authentication keys are defined and valid.

    If the keys are not defined, they are loaded from file by calling
    `loadKeys()'.  If they are expired and no longer valid, the keys are
    rotated by calling `rotateKeys()'.
    """

    if (HMACKey is None) or (AESKey is None):
        loadKeys()

    if (int(time.time()) - creationTime) > const.KEY_ROTATION_TIME:
        rotateKeys()

def decrypt( ticket ):
    """
    Decrypts, verifies and returns the given `ticket'.

    First, the HMAC over the ticket is verified.  If it is valid, the ticket is
    decrypted.  Finally, a `ProtocolState()' object containing the master key
    and the ticket's issue date is returned.  If any of these steps fail,
    `None' is returned.
    """

    assert (ticket is not None) and (len(ticket) == const.TICKET_LENGTH)

    global HMACKey
    global AESKey
    global creationTime

    log.debug("Attempting to decrypt and verify ticket.")

    checkKeys()

    # Verify the ticket's authenticity before decrypting.
    hmac = HMAC.new(HMACKey, ticket[0:80], digestmod=SHA256).digest()
    if hmac != ticket[80:const.TICKET_LENGTH]:
        log.debug("The ticket's HMAC is invalid.  Probably not a ticket.")
        return None

    # Decrypt the ticket to extract the state information.
    aes = AES.new(AESKey, mode=AES.MODE_CBC, IV=ticket[0:IV_LENGTH])
    plainTicket = aes.decrypt(ticket[IV_LENGTH:80])

    issueDate = struct.unpack('I', plainTicket[0:4])[0]
    identifier = plainTicket[4:22]
    masterKey = plainTicket[22:54]

    if not (identifier == IDENTIFIER):
        log.error("The ticket's HMAC is valid but the identifier is invalid.  "
                  "The ticket could be corrupt.")
        return None

    return ProtocolState(masterKey, issueDate=issueDate)

class ProtocolState( object ):

    """
    Defines a ScrambleSuit protocol state contained in a session ticket.

    A protocol state is essentially a master key which can then be used by the
    server to derive session keys.  Besides, a state object contains an issue
    date which specifies the expiry date of a ticket.  This class contains
    methods to check the expiry status of a ticket and to dump it in its raw
    form.
    """

    def __init__( self, masterKey, issueDate=int(time.time()) ):
        """
        The constructor of the `ProtocolState' class.

        The four class variables are initialised.
        """

        self.identifier = IDENTIFIER
        self.masterKey = masterKey
        self.issueDate = issueDate
        # Pad to multiple of 16 bytes to match AES' block size.
        self.pad = "\0\0\0\0\0\0\0\0\0\0"

    def isValid( self ):
        """
        Verifies the expiry date of the object's issue date.

        If the expiry date is not yet reached and the protocol state is still
        valid, `True' is returned.  If the protocol state has expired, `False'
        is returned.
        """

        assert self.issueDate

        lifetime = int(time.time()) - self.issueDate
        if lifetime > const.SESSION_TICKET_LIFETIME:
            log.debug("The ticket is invalid and expired %s ago." %
                      str(datetime.timedelta(seconds=
                      (lifetime - const.SESSION_TICKET_LIFETIME))))
            return False

        log.debug("The ticket is still valid for %s." %
                  str(datetime.timedelta(seconds=
                  (const.SESSION_TICKET_LIFETIME - lifetime))))
        return True

    def __repr__( self ):
        """
        Return a raw string representation of the object's protocol state.

        The length of the returned representation is exactly 64 bytes; a
        multiple of AES' 16-byte block size.  That makes it suitable to be
        encrypted using AES-CBC.
        """

        return struct.pack('I', self.issueDate) + self.identifier + \
                           self.masterKey + self.pad

class SessionTicket( object ):

    """
    Encrypts and authenticates an encapsulated `ProtocolState()' object.

    This class implements a session ticket which can be redeemed by clients.
    The class contains methods to initialise and issue session tickets.
    """

    def __init__( self, masterKey ):
        """
        The constructor of the `SessionTicket()' class.

        The class variables are initialised and the validity of the symmetric
        keys for the session tickets is checked.
        """

        assert (masterKey is not None) and \
               len(masterKey) == const.MASTER_KEY_LENGTH

        checkKeys()

        # Initialisation vector for AES-CBC.
        self.IV = mycrypto.strong_random(IV_LENGTH)

        # The server's (encrypted) protocol state.
        self.state = ProtocolState(masterKey)

        # AES and HMAC keys to encrypt and authenticate the ticket.
        self.symmTicketKey = AESKey
        self.hmacTicketKey = HMACKey

    def issue( self ):
        """
        Returns a ready-to-use session ticket after prior initialisation.

        After the `SessionTicket()' class was initialised with a master key,
        this method encrypts and authenticates the protocol state and returns
        the final result which is ready to be sent over the wire.
        """

        self.state.issueDate = int(time.time())

        # Encrypt the protocol state.
        aes = AES.new(self.symmTicketKey, mode=AES.MODE_CBC, IV=self.IV)
        state = repr(self.state)
        assert (len(state) % AES.block_size) == 0
        cryptedState = aes.encrypt(state)

        # Authenticate the encrypted state and the IV.
        hmac = HMAC.new(self.hmacTicketKey,
                        self.IV + cryptedState, digestmod=SHA256).digest()

        finalTicket = self.IV + cryptedState + hmac
        log.debug("Returning %d-byte ticket." % len(finalTicket))

        return finalTicket

# Alias class name in order to provide a more intuitive API.
new = SessionTicket

# Give ScrambleSuit server operators a way to manually issue new session
# tickets for out-of-band distribution.
if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("ip_addr", type=str, help="The IPv4 address of the "
                        "%s server." % const.TRANSPORT_NAME)
    parser.add_argument("tcp_port", type=int, help="The TCP port of the %s "
                        "server." % const.TRANSPORT_NAME)
    parser.add_argument("ticket_file", type=str, help="The file, the newly "
                        "issued ticket is written to.")
    args = parser.parse_args()

    print "[+] Generating new session ticket."
    masterKey = mycrypto.strong_random(const.MASTER_KEY_LENGTH)
    ticket = SessionTicket(masterKey).issue()

    print "[+] Writing new session ticket to `%s'." % args.ticket_file
    tickets = dict()
    tickets[IPv4Address('TCP', args.ip_addr, args.tcp_port)] = \
           (masterKey, ticket)
    util.writeToFile(pickle.dumps(tickets), args.ticket_file)

    print "[+] Success."
<|endoftext|>


<commit_before>import os
import simplejson as json
import urllib
import contextlib

import logging
log = logging.getLogger(__file__)

def stream_json(u, args=None, extra_path=None):
    """ utility class to stream .json
    """
    ret_val = {}
    url = u
    if extra_path:
        url = "{0}/{1}".format(url, extra_path)
    if args:
        url = "{0}?{1}".format(url, args)
    with contextlib.closing(urllib.urlopen(url)) as stream:
        otp = stream.read()
        ret_val = json.loads(otp)
    return ret_val

def get_json(file_name, path='ott/utils/tests/json', def_val={}):
    """
    utility method to load a static .json file (usually used for testing)
    """
    ret_val = def_val

    # try one with just file_name
    jsn = file_to_json(file_name, None)

    # try two with path prepended to file name
    if jsn is None:
        p = os.path.join(path, file_name)
        jsn = file_to_json(p, None)

    # if we get valid json from either call above, that's the response
    if jsn:
        ret_val = jsn

    return ret_val

def file_to_json(file_path, def_val={}):
    """
    utility method to load a static .json file (usually used in testing)
    """
    ret_val = def_val
    try:
        with open(file_path) as f:
            ret_val = json.load(f)
    except Exception, e:
        log.info(e)

    return ret_val

def str_to_json(str, def_val={}):
    """
    utility function to load a static .json file for mock'ing a service
    """
    ret_val = def_val
    try:
        ret_val = json.loads(str)
    except Exception, e:
        log.info("Couldn't convert {0} to json\n{1}".format(str, e))
    return ret_val

def obj_to_dict(obj):
    """ Represent instance of a class as JSON object (dictionary of named elements)
        :returns a nested dictionary that represents a JSON-encoded object.
        @from: http://stackoverflow.com/a/4682553/2125598
        NOTE: reecursively walks object's hierarchy, creating a nested dict representing the elements of the class
    """
    if(obj is None):
        return None
    if isinstance(obj, (bool, int, long, float, basestring)):
        return obj
    elif isinstance(obj, dict):
        obj = obj.copy()
        for key in obj:
            obj[key.lower()] = obj_to_dict(obj[key])
        return obj
    elif isinstance(obj, list):
        return [obj_to_dict(item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(obj_to_dict([item for item in obj]))
    elif hasattr(obj, '__dict__'):
        return obj_to_dict(obj.__dict__)
    else:
        return repr(obj) # Don't know how to handle, convert to string

def dict_to_json_str(data, pretty_print=False):
    """ dump nested dict into string """
    ret_val = None
    if pretty_print:
        ret_val = json.dumps(data, sort_keys=True, indent=4)
    else:
        ret_val = json.dumps(data)
    return ret_val

def object_to_json_file(file_path, obj, pretty_print=False):
    #import pdb; pdb.set_trace()
    data = obj_to_dict(obj)
    with open(file_path, 'w+') as outfile:
        if pretty_print:
            json.dump(data, outfile, sort_keys=True, indent=4)
        else:
            json.dump(data, outfile)

def json_repr(obj, pretty_print=False):
    ## step 1: call serializer, which walks object tree and returns a cleaned up dict representation of the object
    data = obj_to_dict(obj)

    ## step 2: dump serialized object into json string
    return dict_to_json_str(data, pretty_print)
<commit_msg>minor text fix<commit_after>import os
import simplejson as json
import urllib
import contextlib

import logging
log = logging.getLogger(__file__)

def stream_json(u, args=None, extra_path=None):
    """ utility class to stream .json
    """
    ret_val = {}
    url = u
    if extra_path:
        url = "{0}/{1}".format(url, extra_path)
    if args:
        url = "{0}?{1}".format(url, args)
    with contextlib.closing(urllib.urlopen(url)) as stream:
        data = stream.read()
        ret_val = json.loads(data)
    return ret_val

def get_json(file_name, path='ott/utils/tests/json', def_val={}):
    """
    utility method to load a static .json file (usually used for testing)
    """
    ret_val = def_val

    # try one with just file_name
    jsn = file_to_json(file_name, None)

    # try two with path prepended to file name
    if jsn is None:
        p = os.path.join(path, file_name)
        jsn = file_to_json(p, None)

    # if we get valid json from either call above, that's the response
    if jsn:
        ret_val = jsn

    return ret_val

def file_to_json(file_path, def_val={}):
    """
    utility method to load a static .json file (usually used in testing)
    """
    ret_val = def_val
    try:
        with open(file_path) as f:
            ret_val = json.load(f)
    except Exception, e:
        log.info(e)

    return ret_val

def str_to_json(str, def_val={}):
    """
    utility function to load a static .json file for mock'ing a service
    """
    ret_val = def_val
    try:
        ret_val = json.loads(str)
    except Exception, e:
        log.info("Couldn't convert {0} to json\n{1}".format(str, e))
    return ret_val

def obj_to_dict(obj):
    """ Represent instance of a class as JSON object (dictionary of named elements)
        :returns a nested dictionary that represents a JSON-encoded object.
        @from: http://stackoverflow.com/a/4682553/2125598
        NOTE: reecursively walks object's hierarchy, creating a nested dict representing the elements of the class
    """
    if(obj is None):
        return None
    if isinstance(obj, (bool, int, long, float, basestring)):
        return obj
    elif isinstance(obj, dict):
        obj = obj.copy()
        for key in obj:
            obj[key.lower()] = obj_to_dict(obj[key])
        return obj
    elif isinstance(obj, list):
        return [obj_to_dict(item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(obj_to_dict([item for item in obj]))
    elif hasattr(obj, '__dict__'):
        return obj_to_dict(obj.__dict__)
    else:
        return repr(obj) # Don't know how to handle, convert to string

def dict_to_json_str(data, pretty_print=False):
    """ dump nested dict into string """
    ret_val = None
    if pretty_print:
        ret_val = json.dumps(data, sort_keys=True, indent=4)
    else:
        ret_val = json.dumps(data)
    return ret_val

def object_to_json_file(file_path, obj, pretty_print=False):
    #import pdb; pdb.set_trace()
    data = obj_to_dict(obj)
    with open(file_path, 'w+') as outfile:
        if pretty_print:
            json.dump(data, outfile, sort_keys=True, indent=4)
        else:
            json.dump(data, outfile)

def json_repr(obj, pretty_print=False):
    ## step 1: call serializer, which walks object tree and returns a cleaned up dict representation of the object
    data = obj_to_dict(obj)

    ## step 2: dump serialized object into json string
    return dict_to_json_str(data, pretty_print)
<|endoftext|>


<commit_before>import os, pickle, re
from common.safeprint import safeprint

bountyList = []

class Bounty:
  ip = ""
  btc = ""
  reward = 0
  data = []

  def __init__(self, ipAddress, btcAddress, rewardAmount, dataList=[]):
    self.ip = ipAddress
    self.btc = btcAddress
    self.reward = rewardAmount
    self.data = dataList

  def isValid(self):
    try:
      safeprint("Testing IP address")
      #is IP valid
      b = int(self.ip.split(":")[1]) in range(1024,49152)
      b = int(self.ip.split(":")[0].split(".")[0]) in range(0,256) and b
      b = int(self.ip.split(":")[0].split(".")[1]) in range(0,256) and b
      b = int(self.ip.split(":")[0].split(".")[2]) in range(0,256) and b
      b = int(self.ip.split(":")[0].split(".")[3]) in range(0,256) and b
      if not b:
        return False
      #ping IP
      #is Bitcoin address valid
      safeprint("Testing Bitcoin address")
      address = str(self.btc)
      #The following is a soft check
      #A deeper check will need to be done in order to assure this is correct
      if not re.match(re.compile("^[a-zA-Z1-9]{27,35}$"),address):
        return False
      #is reward valid
      safeprint("Testing reward")
      b = int(self.reward)
      return (b >= 0)
    except:
      return False

  def isPayable(self):
    #check if address has enough
    return False

def verify(string):
  test = pickle.loads(string)
  try:
    safeprint("Testing IP address")
    #is IP valid
    b = int(test.ip.split(":")[1]) in range(1024,49152)
    safeprint(int(test.ip.split(":")[1]))
    safeprint(int(test.ip.split(":")[1]) in range(1024,49152))
    b = int(test.ip.split(":")[0].split(".")[0]) in range(0,256) and b
    b = int(test.ip.split(":")[0].split(".")[1]) in range(0,256) and b
    b = int(test.ip.split(":")[0].split(".")[2]) in range(0,256) and b
    b = int(test.ip.split(":")[0].split(".")[3]) in range(0,256) and b
    #ping IP
    #is Bitcoin address valid
    safeprint("Testing Bitcoin address")
    address = str(test.btc)
    #The following is a soft check
    #A deeper check will need to be done in order to assure this is correct
    if not re.match(re.compile("^[a-zA-Z1-9]{27,35}$"),address):
      return False
    #is reward valid
    safeprint("Testing reward")
    b = int(test.reward)
    return (b >= 0)
  except:
    return False

def saveToFile():
  if os.path.exists("bounties.pickle"):
    pickle.dump(boutyList,"bounties.pickle")
    return True
  return False

def loadFromFile():
  if os.path.exists("settings.conf"):
    bountyList = pickle.load("bounties.pickle")
    return True
  return False

def loadBounties():
  loadFromFile()
  if len(bountyList) is 0:
    requestBounties()
  return len(bountyList) is not 0

def requestBounties(peerList):
  for peer in peerList:
    bountyList.extend(requestBounty(peer))

def requestBounty(peer):
  safeprint("currently unsupported")

def sendBounty(peer):
  safeprint("currently unsupported")
  if len(bountyList) is 0:
    loadBounties()
  #send bounties
  dumpBounties()

def getBounty(charity, factor):
  for bounty in bountyList:
    if best is None:
      best = bounty
    elif best.rewardAmount < bounty.rewardAmount and bounty.isValid() and (isPayable(factor) or charity):
      best = bounty
  return best
<commit_msg>Update bounty.py<commit_after>import os, pickle, re
from common.safeprint import safeprint

bountyList = []

class Bounty:
  ip = ""
  btc = ""
  reward = 0
  data = []

  def __init__(self, ipAddress, btcAddress, rewardAmount, dataList=[]):
    self.ip = ipAddress
    self.btc = btcAddress
    self.reward = rewardAmount
    self.data = dataList

  def isValid(self):
    try:
      safeprint("Testing IP address")
      #is IP valid
      b = int(self.ip.split(":")[1]) in range(1024,49152)
      b = int(self.ip.split(":")[0].split(".")[0]) in range(0,256) and b
      b = int(self.ip.split(":")[0].split(".")[1]) in range(0,256) and b
      b = int(self.ip.split(":")[0].split(".")[2]) in range(0,256) and b
      b = int(self.ip.split(":")[0].split(".")[3]) in range(0,256) and b
      if not b:
        return False
      #ping IP
      #is Bitcoin address valid
      safeprint("Testing Bitcoin address")
      address = str(self.btc)
      #The following is a soft check
      #A deeper check will need to be done in order to assure this is correct
      if not re.match(re.compile("^[a-zA-Z1-9]{27,35}$"),address):
        return False
      #is reward valid
      safeprint("Testing reward")
      b = int(self.reward)
      return (b >= 0)
    except:
      return False

  def isPayable(self):
    #check if address has enough
    return False

def verify(string):
  test = pickle.loads(string)
  try:
    safeprint("Testing IP address")
    #is IP valid
    b = int(test.ip.split(":")[1]) in range(1024,49152)
    safeprint(int(test.ip.split(":")[1]))
    safeprint(int(test.ip.split(":")[1]) in range(1024,49152))
    safeprint(b)
    safeprint(int(test.ip.split(":")[0].split(".")[0]) in range(0,256) and b)
    safeprint(b and int(test.ip.split(":")[0].split(".")[0]) in range(0,256))
    b = int(test.ip.split(":")[0].split(".")[0]) in range(0,256) and b
    b = int(test.ip.split(":")[0].split(".")[1]) in range(0,256) and b
    b = int(test.ip.split(":")[0].split(".")[2]) in range(0,256) and b
    b = int(test.ip.split(":")[0].split(".")[3]) in range(0,256) and b
    #ping IP
    #is Bitcoin address valid
    safeprint("Testing Bitcoin address")
    address = str(test.btc)
    #The following is a soft check
    #A deeper check will need to be done in order to assure this is correct
    if not re.match(re.compile("^[a-zA-Z1-9]{27,35}$"),address):
      return False
    #is reward valid
    safeprint("Testing reward")
    b = int(test.reward)
    return (b >= 0)
  except:
    return False

def saveToFile():
  if os.path.exists("bounties.pickle"):
    pickle.dump(boutyList,"bounties.pickle")
    return True
  return False

def loadFromFile():
  if os.path.exists("settings.conf"):
    bountyList = pickle.load("bounties.pickle")
    return True
  return False

def loadBounties():
  loadFromFile()
  if len(bountyList) is 0:
    requestBounties()
  return len(bountyList) is not 0

def requestBounties(peerList):
  for peer in peerList:
    bountyList.extend(requestBounty(peer))

def requestBounty(peer):
  safeprint("currently unsupported")

def sendBounty(peer):
  safeprint("currently unsupported")
  if len(bountyList) is 0:
    loadBounties()
  #send bounties
  dumpBounties()

def getBounty(charity, factor):
  for bounty in bountyList:
    if best is None:
      best = bounty
    elif best.rewardAmount < bounty.rewardAmount and bounty.isValid() and (isPayable(factor) or charity):
      best = bounty
  return best
<|endoftext|>