<commit_before>import os
import Queue
import json
import hashlib
from datetime import datetime

import sublime
import sublime_plugin
import dmp_monkey
dmp_monkey.monkey_patch()
from lib import diff_match_patch as dmp

import msg
import shared as G
import utils

MODIFIED_EVENTS = Queue.Queue()
BUFS = {}

settings = sublime.load_settings('Floobits.sublime-settings')

def get_text(view):
    return view.substr(sublime.Region(0, view.size()))

def get_view(buf_id):
    buf = BUFS[buf_id]
    for view in G.ROOM_WINDOW.views():
        if not view.file_name():
            continue
        if buf['path'] == utils.to_rel_path(view.file_name()):
            return view
    return None

def create_view(buf):
    path = utils.get_full_path(buf['path'])
    view = G.ROOM_WINDOW.open_file(path)
    if view:
        msg.debug('Created view', view.name() or view.file_name())
    return view

def get_buf(view):
    if view.is_scratch():
        return None
    if not view.file_name():
        return None
    rel_path = utils.to_rel_path(view.file_name())
    for buf_id, buf in BUFS.iteritems():
        if rel_path == buf['path']:
            return buf
    return None

def save_buf(buf):
    path = utils.get_full_path(buf['path'])
    utils.mkdir(os.path.split(path)[0])
    with open(path, 'wb') as fd:
        fd.write(buf['buf'].encode('utf-8'))

def delete_buf(buf_id):
    # TODO: somehow tell the user about this. maybe delete on disk too?
    del BUFS[buf_id]

class FlooPatch(object):

    def __init__(self, view, buf):
        self.buf = buf
        self.view = view
        self.current = get_text(view)
        self.previous = buf['buf']
        self.md5_before = hashlib.md5(self.previous.encode('utf-8')).hexdigest()

    def __str__(self):
        return '%s - %s - %s' % (self.buf['id'], self.buf['path'], self.view.buffer_id())

    def patches(self):
        return dmp.diff_match_patch().patch_make(self.previous, self.current)

    def to_json(self):
        patches = self.patches()
        if len(patches) == 0:
            return None
        msg.debug('sending %s patches' % len(patches))
        patch_str = ''
        for patch in patches:
            patch_str += str(patch)
        return json.dumps({
            'id': self.buf['id'],
            'md5_after': hashlib.md5(self.current.encode('utf-8')).hexdigest(),
            'md5_before': self.md5_before,
            'path': self.buf['path'],
            'patch': patch_str,
            'name': 'patch'
        })

class Listener(sublime_plugin.EventListener):
    views_changed = []
    selection_changed = []
    agent = None

    def __init__(self, *args, **kwargs):
        sublime_plugin.EventListener.__init__(self, *args, **kwargs)
        self.between_save_events = {}

    @staticmethod
    def set_agent(agent):
        Listener.agent = agent

    @staticmethod
    def push():
        reported = set()
        while Listener.views_changed:
            view, buf = Listener.views_changed.pop()

            vb_id = view.buffer_id()
            if vb_id in reported:
                continue
            if 'buf' not in buf:
                msg.debug('No data for buf %s %s yet. Skipping sending patch' % (buf['id'], buf['path']))
                continue

            reported.add(vb_id)
            patch = FlooPatch(view, buf)
            # update the current copy of the buffer
            buf['buf'] = patch.current
            buf['md5'] = hashlib.md5(patch.current.encode('utf-8')).hexdigest()
            if Listener.agent:
                Listener.agent.put(patch.to_json())
            else:
                msg.debug('Not connected. Discarding view change.')

        while Listener.selection_changed:
            view, buf = Listener.selection_changed.pop()
            vb_id = view.buffer_id()
            if vb_id in reported:
                continue

            reported.add(vb_id)
            sel = view.sel()
            highlight_json = json.dumps({
                'id': buf['id'],
                'name': 'highlight',
                'ranges': [[x.a, x.b] for x in sel]
            })
            if Listener.agent:
                Listener.agent.put(highlight_json)
            else:
                msg.debug('Not connected. Discarding selection change.')

        sublime.set_timeout(Listener.push, 100)

    @staticmethod
    def apply_patch(patch_data):
        buf_id = patch_data['id']
        buf = BUFS[buf_id]
        view = get_view(buf_id)
        DMP = dmp.diff_match_patch()
        if len(patch_data['patch']) == 0:
            msg.error('wtf? no patches to apply. server is being stupid')
            return
        msg.debug('patch is', patch_data['patch'])
        dmp_patches = DMP.patch_fromText(patch_data['patch'])
        # TODO: run this in a separate thread
        if view:
            old_text = get_text(view)
        else:
            old_text = buf.get('buf', '')
        md5_before = hashlib.md5(old_text.encode('utf-8')).hexdigest()
        if md5_before != patch_data['md5_before']:
            msg.warn('starting md5s don\'t match for %s. this is dangerous!' % buf['path'])

        t = DMP.patch_apply(dmp_patches, old_text)

        clean_patch = True
        for applied_patch in t[1]:
            if not applied_patch:
                clean_patch = False
                break

        if G.DEBUG:
            if len(t[0]) == 0:
                msg.debug('OMG EMPTY!')
                msg.debug('Starting data:', buf['buf'])
                msg.debug('Patch:', patch_data['patch'])
            if '\x01' in t[0]:
                msg.debug('FOUND CRAZY BYTE IN BUFFER')
                msg.debug('Starting data:', buf['buf'])
                msg.debug('Patch:', patch_data['patch'])

        if not clean_patch:
            msg.error('failed to patch %s cleanly. re-fetching buffer' % buf['path'])
            return Listener.get_buf(buf_id)

        cur_hash = hashlib.md5(t[0].encode('utf-8')).hexdigest()
        if cur_hash != patch_data['md5_after']:
            msg.warn(
                '%s new hash %s != expected %s. re-fetching buffer...' %
                (buf['path'], cur_hash, patch_data['md5_after'])
            )
            return Listener.get_buf(buf_id)

        buf['buf'] = t[0].decode('utf-8')
        buf['md5'] = cur_hash

        if not view:
            save_buf(buf)
            return

        selections = [x for x in view.sel()]  # deep copy
        regions = []
        for patch in t[2]:
            offset = patch[0]
            length = patch[1]
            patch_text = patch[2]
            region = sublime.Region(offset, offset + length)
            regions.append(region)
            MODIFIED_EVENTS.put(1)
            try:
                edit = view.begin_edit()
                view.replace(edit, region, patch_text.decode('utf-8'))
            except:
                raise
            else:
                new_sels = []
                for sel in selections:
                    a = sel.a
                    b = sel.b
                    new_offset = len(patch_text) - length
                    if sel.a > offset:
                        a += new_offset
                    if sel.b > offset:
                        b += new_offset
                    new_sels.append(sublime.Region(a, b))
                selections = [x for x in new_sels]
            finally:
                view.end_edit(edit)
        view.sel().clear()
        region_key = 'floobits-patch-' + patch_data['username']
        view.add_regions(region_key, regions, 'floobits.patch', 'circle', sublime.DRAW_OUTLINED)
        sublime.set_timeout(lambda: view.erase_regions(region_key), 1000)
        for sel in selections:
            view.sel().add(sel)

        now = datetime.now()
        view.set_status('Floobits', 'Changed by %s at %s' % (patch_data['username'], now.strftime('%H:%M')))

    @staticmethod
    def get_buf(buf_id):
        req = {
            'name': 'get_buf',
            'id': buf_id
        }
        Listener.agent.put(json.dumps(req))

    @staticmethod
    def create_buf(path):
        if not utils.is_shared(path):
            msg.error('Skipping adding %s because it is not in shared path %s.' % (path, G.PROJECT_PATH))
            return
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                # Don't care about hidden stuff
                dirnames[:] = [d for d in dirnames if d[0] != '.']
                for f in filenames:
                    f_path = os.path.join(dirpath, f)
                    if f[0] == '.':
                        msg.log('Not creating buf for hidden file %s' % f_path)
                    else:
                        Listener.create_buf(f_path)
            return
        try:
            buf_fd = open(path, 'rb')
            buf = buf_fd.read().decode('utf-8')
            rel_path = utils.to_rel_path(path)
            msg.log('creating buffer ', rel_path)
            event = {
                'name': 'create_buf',
                'buf': buf,
                'path': rel_path,
            }
            Listener.agent.put(json.dumps(event))
        except (IOError, OSError):
            msg.error('Failed to open %s.' % path)
        except Exception as e:
            msg.error('Failed to create buffer %s: %s' % (path, str(e)))

    @staticmethod
    def delete_buf(path):
        if not utils.is_shared(path):
            msg.error('Skipping deleting %s because it is not in shared path %s.' % (path, G.PROJECT_PATH))
            return
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                # Don't care about hidden stuff
                dirnames[:] = [d for d in dirnames if d[0] != '.']
                for f in filenames:
                    f_path = os.path.join(dirpath, f)
                    if f[0] == '.':
                        msg.log('Not deleting buf for hidden file %s' % f_path)
                    else:
                        Listener.delete_buf(f_path)
            return
        buf_to_delete = None
        rel_path = utils.to_rel_path(path)
        for buf_id, buf in BUFS.items():
            if rel_path == buf['path']:
                buf_to_delete = buf
                break
        if buf_to_delete is None:
            msg.error('%s is not in this room' % path)
            return
        msg.log('deleting buffer ', rel_path)
        event = {
            'name': 'delete_buf',
            'id': buf_to_delete['id'],
        }
        Listener.agent.put(json.dumps(event))

    @staticmethod
    def update_view(buf, view=None):
        view = view or get_view(buf['id'])
        visible_region = view.visible_region()
        viewport_position = view.viewport_position()
        region = sublime.Region(0, view.size())
        selections = [x for x in view.sel()]  # deep copy
        MODIFIED_EVENTS.put(1)
        try:
            edit = view.begin_edit()
            view.replace(edit, region, buf['buf'])
        except Exception as e:
            msg.error('Exception updating view: %s' % e)
        finally:
            view.end_edit(edit)
        sublime.set_timeout(lambda: view.set_viewport_position(viewport_position, False), 0)
        view.sel().clear()
        view.show(visible_region, False)
        for sel in selections:
            view.sel().add(sel)
        view.set_read_only(G.READ_ONLY)
        if G.READ_ONLY:
            view.set_status('Floobits', 'You don\'t have write permission. Buffer is read-only.')

    @staticmethod
    def highlight(buf_id, region_key, username, ranges):
        view = get_view(buf_id)
        if not view:
            return
        regions = []
        for r in ranges:
            regions.append(sublime.Region(*r))
        view.erase_regions(region_key)
        view.add_regions(region_key, regions, region_key, 'dot', sublime.DRAW_OUTLINED)

    def id(self, view):
        return view.buffer_id()

    def name(self, view):
        return view.file_name()

    def on_new(self, view):
        msg.debug('new', self.name(view))

    def on_clone(self, view):
        msg.debug('clone', self.name(view))

    def on_load(self, view):
        msg.debug('load', self.name(view))

    def on_pre_save(self, view):
        p = view.name()
        if view.file_name():
            p = utils.to_rel_path(view.file_name())
        self.between_save_events[view.buffer_id()] = p

    def on_post_save(self, view):
        def cleanup():
            del self.between_save_events[view.buffer_id()]
        if view == G.CHAT_VIEW or view.file_name() == G.CHAT_VIEW_PATH:
            return cleanup()
        else:
            print G.CHAT_VIEW_PATH, "not", view.file_name()
        event = None
        buf = get_buf(view)
        name = utils.to_rel_path(view.file_name())
        old_name = self.between_save_events[view.buffer_id()]

        if buf is None:
            if utils.is_shared(view.file_name()):
                msg.log('new buffer ', name, view.file_name())
                event = {
                    'name': 'create_buf',
                    'buf': get_text(view),
                    'path': name
                }
        elif name != old_name:
            if utils.is_shared(view.file_name()):
                msg.log('renamed buffer {0} to {1}'.format(old_name, name))
                event = {
                    'name': 'rename_buf',
                    'id': buf['id'],
                    'path': name
                }
            else:
                msg.log('deleting buffer from shared: {0}'.format(name))
                event = {
                    'name': 'delete_buf',
                    'id': buf['id'],
                }

        if event and Listener.agent:
            Listener.agent.put(json.dumps(event))

        cleanup()

    def on_modified(self, view):
        try:
            MODIFIED_EVENTS.get_nowait()
        except Queue.Empty:
            self.add(view)
        else:
            MODIFIED_EVENTS.task_done()

    def on_selection_modified(self, view):
        buf = get_buf(view)
        if buf:
            msg.debug('selection in view %s is buf id %s' % (buf['path'], buf['id']))
            self.selection_changed.append((view, buf))

    def on_activated(self, view):
        self.add(view)

    def add(self, view):
        buf = get_buf(view)
        if buf:
            msg.debug('changed view %s buf id %s' % (buf['path'], buf['id']))
            self.views_changed.append((view, buf))
<commit_msg>Fix some stupid unicode errors.<commit_after>import os
import Queue
import json
import hashlib
from datetime import datetime

import sublime
import sublime_plugin
import dmp_monkey
dmp_monkey.monkey_patch()
from lib import diff_match_patch as dmp

import msg
import shared as G
import utils

MODIFIED_EVENTS = Queue.Queue()
BUFS = {}

settings = sublime.load_settings('Floobits.sublime-settings')

def get_text(view):
    return view.substr(sublime.Region(0, view.size()))

def get_view(buf_id):
    buf = BUFS[buf_id]
    for view in G.ROOM_WINDOW.views():
        if not view.file_name():
            continue
        if buf['path'] == utils.to_rel_path(view.file_name()):
            return view
    return None

def create_view(buf):
    path = utils.get_full_path(buf['path'])
    view = G.ROOM_WINDOW.open_file(path)
    if view:
        msg.debug('Created view', view.name() or view.file_name())
    return view

def get_buf(view):
    if view.is_scratch():
        return None
    if not view.file_name():
        return None
    rel_path = utils.to_rel_path(view.file_name())
    for buf_id, buf in BUFS.iteritems():
        if rel_path == buf['path']:
            return buf
    return None

def save_buf(buf):
    path = utils.get_full_path(buf['path'])
    utils.mkdir(os.path.split(path)[0])
    with open(path, 'wb') as fd:
        fd.write(buf['buf'].encode('utf-8'))

def delete_buf(buf_id):
    # TODO: somehow tell the user about this. maybe delete on disk too?
    del BUFS[buf_id]

class FlooPatch(object):

    def __init__(self, view, buf):
        self.buf = buf
        self.view = view
        self.current = get_text(view)
        self.previous = buf['buf']
        self.md5_before = hashlib.md5(self.previous.encode('utf-8')).hexdigest()

    def __str__(self):
        return '%s - %s - %s' % (self.buf['id'], self.buf['path'], self.view.buffer_id())

    def patches(self):
        return dmp.diff_match_patch().patch_make(self.previous, self.current)

    def to_json(self):
        patches = self.patches()
        if len(patches) == 0:
            return None
        msg.debug('sending %s patches' % len(patches))
        patch_str = ''
        for patch in patches:
            patch_str += str(patch)
        return json.dumps({
            'id': self.buf['id'],
            'md5_after': hashlib.md5(self.current.encode('utf-8')).hexdigest(),
            'md5_before': self.md5_before,
            'path': self.buf['path'],
            'patch': patch_str,
            'name': 'patch'
        })

class Listener(sublime_plugin.EventListener):
    views_changed = []
    selection_changed = []
    agent = None

    def __init__(self, *args, **kwargs):
        sublime_plugin.EventListener.__init__(self, *args, **kwargs)
        self.between_save_events = {}

    @staticmethod
    def set_agent(agent):
        Listener.agent = agent

    @staticmethod
    def push():
        reported = set()
        while Listener.views_changed:
            view, buf = Listener.views_changed.pop()

            vb_id = view.buffer_id()
            if vb_id in reported:
                continue
            if 'buf' not in buf:
                msg.debug('No data for buf %s %s yet. Skipping sending patch' % (buf['id'], buf['path']))
                continue

            reported.add(vb_id)
            patch = FlooPatch(view, buf)
            # update the current copy of the buffer
            buf['buf'] = patch.current
            buf['md5'] = hashlib.md5(patch.current.encode('utf-8')).hexdigest()
            if Listener.agent:
                Listener.agent.put(patch.to_json())
            else:
                msg.debug('Not connected. Discarding view change.')

        while Listener.selection_changed:
            view, buf = Listener.selection_changed.pop()
            vb_id = view.buffer_id()
            if vb_id in reported:
                continue

            reported.add(vb_id)
            sel = view.sel()
            highlight_json = json.dumps({
                'id': buf['id'],
                'name': 'highlight',
                'ranges': [[x.a, x.b] for x in sel]
            })
            if Listener.agent:
                Listener.agent.put(highlight_json)
            else:
                msg.debug('Not connected. Discarding selection change.')

        sublime.set_timeout(Listener.push, 100)

    @staticmethod
    def apply_patch(patch_data):
        buf_id = patch_data['id']
        buf = BUFS[buf_id]
        view = get_view(buf_id)
        DMP = dmp.diff_match_patch()
        if len(patch_data['patch']) == 0:
            msg.error('wtf? no patches to apply. server is being stupid')
            return
        msg.debug('patch is', patch_data['patch'])
        dmp_patches = DMP.patch_fromText(patch_data['patch'])
        # TODO: run this in a separate thread
        if view:
            old_text = get_text(view)
        else:
            old_text = buf.get('buf', '')
        md5_before = hashlib.md5(old_text.encode('utf-8')).hexdigest()
        if md5_before != patch_data['md5_before']:
            msg.warn('starting md5s don\'t match for %s. this is dangerous!' % buf['path'])

        t = DMP.patch_apply(dmp_patches, old_text)

        clean_patch = True
        for applied_patch in t[1]:
            if not applied_patch:
                clean_patch = False
                break

        if G.DEBUG:
            if len(t[0]) == 0:
                msg.debug('OMG EMPTY!')
                msg.debug('Starting data:', buf['buf'])
                msg.debug('Patch:', patch_data['patch'])
            if '\x01' in t[0]:
                msg.debug('FOUND CRAZY BYTE IN BUFFER')
                msg.debug('Starting data:', buf['buf'])
                msg.debug('Patch:', patch_data['patch'])

        if not clean_patch:
            msg.error('failed to patch %s cleanly. re-fetching buffer' % buf['path'])
            return Listener.get_buf(buf_id)

        cur_hash = hashlib.md5(t[0].encode('utf-8')).hexdigest()
        if cur_hash != patch_data['md5_after']:
            msg.warn(
                '%s new hash %s != expected %s. re-fetching buffer...' %
                (buf['path'], cur_hash, patch_data['md5_after'])
            )
            return Listener.get_buf(buf_id)

        buf['buf'] = t[0]
        buf['md5'] = cur_hash

        if not view:
            save_buf(buf)
            return

        selections = [x for x in view.sel()]  # deep copy
        regions = []
        for patch in t[2]:
            offset = patch[0]
            length = patch[1]
            patch_text = patch[2]
            region = sublime.Region(offset, offset + length)
            regions.append(region)
            MODIFIED_EVENTS.put(1)
            try:
                edit = view.begin_edit()
                view.replace(edit, region, patch_text)
            except:
                raise
            else:
                new_sels = []
                for sel in selections:
                    a = sel.a
                    b = sel.b
                    new_offset = len(patch_text) - length
                    if sel.a > offset:
                        a += new_offset
                    if sel.b > offset:
                        b += new_offset
                    new_sels.append(sublime.Region(a, b))
                selections = [x for x in new_sels]
            finally:
                view.end_edit(edit)
        view.sel().clear()
        region_key = 'floobits-patch-' + patch_data['username']
        view.add_regions(region_key, regions, 'floobits.patch', 'circle', sublime.DRAW_OUTLINED)
        sublime.set_timeout(lambda: view.erase_regions(region_key), 1000)
        for sel in selections:
            view.sel().add(sel)

        now = datetime.now()
        view.set_status('Floobits', 'Changed by %s at %s' % (patch_data['username'], now.strftime('%H:%M')))

    @staticmethod
    def get_buf(buf_id):
        req = {
            'name': 'get_buf',
            'id': buf_id
        }
        Listener.agent.put(json.dumps(req))

    @staticmethod
    def create_buf(path):
        if not utils.is_shared(path):
            msg.error('Skipping adding %s because it is not in shared path %s.' % (path, G.PROJECT_PATH))
            return
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                # Don't care about hidden stuff
                dirnames[:] = [d for d in dirnames if d[0] != '.']
                for f in filenames:
                    f_path = os.path.join(dirpath, f)
                    if f[0] == '.':
                        msg.log('Not creating buf for hidden file %s' % f_path)
                    else:
                        Listener.create_buf(f_path)
            return
        try:
            buf_fd = open(path, 'rb')
            buf = buf_fd.read().decode('utf-8')
            rel_path = utils.to_rel_path(path)
            msg.log('creating buffer ', rel_path)
            event = {
                'name': 'create_buf',
                'buf': buf,
                'path': rel_path,
            }
            Listener.agent.put(json.dumps(event))
        except (IOError, OSError):
            msg.error('Failed to open %s.' % path)
        except Exception as e:
            msg.error('Failed to create buffer %s: %s' % (path, str(e)))

    @staticmethod
    def delete_buf(path):
        if not utils.is_shared(path):
            msg.error('Skipping deleting %s because it is not in shared path %s.' % (path, G.PROJECT_PATH))
            return
        if os.path.isdir(path):
            for dirpath, dirnames, filenames in os.walk(path):
                # Don't care about hidden stuff
                dirnames[:] = [d for d in dirnames if d[0] != '.']
                for f in filenames:
                    f_path = os.path.join(dirpath, f)
                    if f[0] == '.':
                        msg.log('Not deleting buf for hidden file %s' % f_path)
                    else:
                        Listener.delete_buf(f_path)
            return
        buf_to_delete = None
        rel_path = utils.to_rel_path(path)
        for buf_id, buf in BUFS.items():
            if rel_path == buf['path']:
                buf_to_delete = buf
                break
        if buf_to_delete is None:
            msg.error('%s is not in this room' % path)
            return
        msg.log('deleting buffer ', rel_path)
        event = {
            'name': 'delete_buf',
            'id': buf_to_delete['id'],
        }
        Listener.agent.put(json.dumps(event))

    @staticmethod
    def update_view(buf, view=None):
        view = view or get_view(buf['id'])
        visible_region = view.visible_region()
        viewport_position = view.viewport_position()
        region = sublime.Region(0, view.size())
        selections = [x for x in view.sel()]  # deep copy
        MODIFIED_EVENTS.put(1)
        try:
            edit = view.begin_edit()
            view.replace(edit, region, buf['buf'])
        except Exception as e:
            msg.error('Exception updating view: %s' % e)
        finally:
            view.end_edit(edit)
        sublime.set_timeout(lambda: view.set_viewport_position(viewport_position, False), 0)
        view.sel().clear()
        view.show(visible_region, False)
        for sel in selections:
            view.sel().add(sel)
        view.set_read_only(G.READ_ONLY)
        if G.READ_ONLY:
            view.set_status('Floobits', 'You don\'t have write permission. Buffer is read-only.')

    @staticmethod
    def highlight(buf_id, region_key, username, ranges):
        view = get_view(buf_id)
        if not view:
            return
        regions = []
        for r in ranges:
            regions.append(sublime.Region(*r))
        view.erase_regions(region_key)
        view.add_regions(region_key, regions, region_key, 'dot', sublime.DRAW_OUTLINED)

    def id(self, view):
        return view.buffer_id()

    def name(self, view):
        return view.file_name()

    def on_new(self, view):
        msg.debug('new', self.name(view))

    def on_clone(self, view):
        msg.debug('clone', self.name(view))

    def on_load(self, view):
        msg.debug('load', self.name(view))

    def on_pre_save(self, view):
        p = view.name()
        if view.file_name():
            p = utils.to_rel_path(view.file_name())
        self.between_save_events[view.buffer_id()] = p

    def on_post_save(self, view):
        def cleanup():
            del self.between_save_events[view.buffer_id()]
        if view == G.CHAT_VIEW or view.file_name() == G.CHAT_VIEW_PATH:
            return cleanup()
        else:
            print G.CHAT_VIEW_PATH, "not", view.file_name()
        event = None
        buf = get_buf(view)
        name = utils.to_rel_path(view.file_name())
        old_name = self.between_save_events[view.buffer_id()]

        if buf is None:
            if utils.is_shared(view.file_name()):
                msg.log('new buffer ', name, view.file_name())
                event = {
                    'name': 'create_buf',
                    'buf': get_text(view),
                    'path': name
                }
        elif name != old_name:
            if utils.is_shared(view.file_name()):
                msg.log('renamed buffer {0} to {1}'.format(old_name, name))
                event = {
                    'name': 'rename_buf',
                    'id': buf['id'],
                    'path': name
                }
            else:
                msg.log('deleting buffer from shared: {0}'.format(name))
                event = {
                    'name': 'delete_buf',
                    'id': buf['id'],
                }

        if event and Listener.agent:
            Listener.agent.put(json.dumps(event))

        cleanup()

    def on_modified(self, view):
        try:
            MODIFIED_EVENTS.get_nowait()
        except Queue.Empty:
            self.add(view)
        else:
            MODIFIED_EVENTS.task_done()

    def on_selection_modified(self, view):
        buf = get_buf(view)
        if buf:
            msg.debug('selection in view %s is buf id %s' % (buf['path'], buf['id']))
            self.selection_changed.append((view, buf))

    def on_activated(self, view):
        self.add(view)

    def add(self, view):
        buf = get_buf(view)
        if buf:
            msg.debug('changed view %s buf id %s' % (buf['path'], buf['id']))
            self.views_changed.append((view, buf))
<|endoftext|>


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

""" Assumes fabric environment already set up """

__author__ = 'ilyash'

import subprocess
import sys
import time
import os

from fabric import operations
import fabric.api
from path import path
import sh

from cosmo_tester.framework.testenv import TestCase
from cosmo_tester.framework.util import YamlPatcher

CHEF_SERVER_COOKBOOK_ZIP_URL = (
    'https://github.com/opscode-cookbooks/chef-server/archive/'
    'c588a4c401d3fac14f70d3285fe49eb4dccd9759.zip'
)

CHEF_SERVER_COOKBOOKS_TAR_GZS = (
    (
        'create-file',
        'https://github.com/ilyash/cookbook-create-file/archive/'
        'bd8218a6a9b3e33a165042241c50908ccb6145d1.tar.gz'
    ),
)

KNIFE_PARAMS = '-u admin -k ~/admin.pem'

IMAGE_NAME = 'Ubuntu-NP'
FLAVOR_NAME = 'm1.small'

def _use_cookbook(cookbook_name, cookbook_url):
    """ Downloads cookbook from given url and uploads it to the Chef server """
    fabric.api.run('mkdir -p ~/cookbooks/{0}'.format(cookbook_name))
    # Next line was inspired by (sorry, flake8)
    # https://github.com/opscode-cookbooks/chef-server/blame/
    # c588a4c401d3fac14f70d3285fe49eb4dccd9759/README.md#L158
    fabric.api.run('wget -qO- {0} | tar xvzC ~/cookbooks/{1}'
                   ' --strip-components=1'.format(cookbook_url, cookbook_name))
    fabric.api.run('knife cookbook upload {0} --cookbook-path ~/cookbooks {1}'
                   .format(KNIFE_PARAMS, cookbook_name))
    fabric.api.run('knife cookbook list {0} | grep -F {1}'
                   .format(KNIFE_PARAMS, cookbook_name))

def _userize_file(original_path):
    """ Places the file under user's home directory and make it
        permissions-wise accessible """
    fabric.api.sudo("cp -a {path} ~{user}/ && chown {user} ~{user}/{basename}"
                    .format(path=original_path,
                            basename=str(path(original_path).basename()),
                            user=fabric.api.env['user']))

def setup_chef_server(local_dir, cookbooks):
    _userize_file("/etc/chef-server/admin.pem")
    for cb in cookbooks:
        _use_cookbook(*cb)

    _userize_file("/etc/chef-server/chef-validator.pem")
    operations.get('~/chef-validator.pem', str(local_dir))

def find_node_state(node_name, nodes_state):
    pfx = node_name + '_'
    matches = [v for k, v in nodes_state.items() if k.startswith(pfx)]
    if len(matches) != 1:
        raise RuntimeError("Failed to find node {0}".format(node_name))
    return matches[0]

def get_nodes_of_type(blueprint, type_):
    return [n for n in blueprint.obj['blueprint']['nodes']
            if n['type'] == type_]

def get_agent_key_file(env):
    agent_key_file = path(os.path.expanduser(env.agent_key_path)).abspath()
    if not agent_key_file.exists():
        raise RuntimeError("Agent key file {0} does not exist".format(
            agent_key_file))
    return agent_key_file

def update_blueprint(env, blueprint, hostname, userdata_vars=None):
    hostname_base = 'system-test-{0}-{1}'.format(
        time.strftime("%Y%m%d-%H%M"), hostname)
    agent_key_file = get_agent_key_file(env)
    vms = get_nodes_of_type(blueprint, 'cloudify.openstack.server')
    if len(vms) > 1:
        hostnames = ['{0}-{1:2}'.format(hostname_base, i)
                     for i in range(0, len(vms))]
    else:
        hostnames = [hostname_base]

    users = []
    for vm_idx, vm in enumerate(vms):
        vm_hostname = hostnames[vm_idx]
        vm['properties']['worker_config']['key'] = (
            '~/.ssh/' + str(agent_key_file.basename()))

        # vm['properties']['server'] does not exist when using existing one
        if 'server' in vm['properties']:
            vm['properties']['server'].update({
                'flavor_name': FLAVOR_NAME,
                'image_name': IMAGE_NAME,
                'key_name': env.agent_keypair_name,
                'name': vm_hostname,
            })
            vm['properties']['management_network_name'] = (
                env.management_network_name)
            vm['properties']['server']['security_groups'].append(
                env.agents_security_group)
            props = vm['properties']['server']
            if 'userdata' in props:
                props['userdata'] = props['userdata'].format(
                    hostname=vm_hostname, **(userdata_vars or {}))
        users.append(vm['properties']['worker_config']['user'])

    fips = get_nodes_of_type(blueprint, 'cloudify.openstack.floatingip')
    for fip in fips:
            fip_fip = fip['properties']['floatingip']
            fip_fip['floating_network_name'] = env.external_network_name

    return {'hostnames': hostnames, 'users': users}

class ChefPluginClientTest(TestCase):

    def setUp(self, *args, **kwargs):

        super(ChefPluginClientTest, self).setUp(*args, **kwargs)
        agent_key_file = get_agent_key_file(self.env)

        blueprint_dir = self.copy_blueprint('chef-plugin')
        self.blueprint_yaml = blueprint_dir / 'chef-server-by-chef-solo.yaml'

        with YamlPatcher(self.blueprint_yaml) as blueprint:
            bp_info = update_blueprint(self.env, blueprint, 'chef-server')

        self.chef_server_hostname = bp_info['hostnames'][0]

        cookbooks_dir = blueprint_dir / 'cookbooks'
        cookbooks_dir.mkdir()

        def run(*args, **kwargs):
            return subprocess.check_output(*args, stderr=sys.stderr, **kwargs)

        with cookbooks_dir:
            run([
                'wget', '-q', '-O', 'chef-server.zip',
                CHEF_SERVER_COOKBOOK_ZIP_URL,
                ])
            run(['unzip', 'chef-server.zip'])
            chef_cookbook_dir = cookbooks_dir.glob('chef-server-*')[0]
            run(['mv', chef_cookbook_dir, 'chef-server'])
            # Next line because Chef cookbooks are required
            # to declare all dependencies, even if they don't use them.
            # We don't need git, it's only used in chef-cookbook::dev recipe.
            run(['sed', '-i', "/depends 'git'/d", 'chef-server/metadata.rb'])

        with blueprint_dir:
            run(['tar', 'czf', 'cookbooks.tar.gz', 'cookbooks'])

        id_ = self.test_id + '-chef-server'
        before, after = self.upload_deploy_and_execute_install(id_, id_)

        fip_node = find_node_state('ip', after['node_state'][id_])
        self.chef_server_ip = fip_node['runtimeInfo']['floating_ip_address']

        fabric_env = fabric.api.env
        fabric_env.update({
            'timeout': 30,
            'user': bp_info['users'][0],
            'key_filename': str(agent_key_file),
            'host_string': self.chef_server_ip,
        })

        setup_chef_server(blueprint_dir, CHEF_SERVER_COOKBOOKS_TAR_GZS)
        self.blueprint_dir = blueprint_dir

    def test_chef_client(self):
        blueprint_dir = self.blueprint_dir
        self.blueprint_yaml = blueprint_dir / 'chef-client-test.yaml'
        with YamlPatcher(self.blueprint_yaml) as blueprint:
            update_blueprint(self.env, blueprint, 'chef-server', {
                'chef_server_ip': self.chef_server_ip,
                'chef_server_hostname': self.chef_server_hostname,
            })
            chef_node = get_nodes_of_type(blueprint,
                                          'cloudify.types.chef.db_server')[0]
            chef_config = chef_node['properties']['chef_config']
            chef_config['chef_server_url'] = 'https://{0}:443'.format(
                self.chef_server_ip)
            chef_config['validation_client_name'] = 'chef-validator'
            chef_config['validation_key'] = (
                path(blueprint_dir) / 'chef-validator.pem').text()

        id_ = self.test_id + '-chef-client-' + str(int(time.time()))
        before, after = self.upload_deploy_and_execute_install(id_, id_)

        fip_node = find_node_state('ip', after['node_state'][id_])
        chef_client_ip = fip_node['runtimeInfo']['floating_ip_address']

        fabric_env = fabric.api.env
        fabric_env.update({
            # XXX: sometime - same user for connection is accidental
            #      todo: replace it with update_blueprint()'s bp_info,
            #            as in setUp()
            'host_string': chef_client_ip,
        })

        out = fabric.api.run('cat /tmp/blueprint.txt')
        self.assertEquals(out, 'Great success!')

class ChefPluginSoloTest(TestCase):

    def setUp(self, *args, **kwargs):

        super(ChefPluginSoloTest, self).setUp(*args, **kwargs)

        self.blueprint_dir = self.copy_blueprint('chef-plugin')

        # Get resources
        with self.blueprint_dir:
            # Cookbooks
            for name, url in CHEF_SERVER_COOKBOOKS_TAR_GZS:
                sh.mkdir('-p', 'cookbooks/' + name)
                sh.tar(sh.wget('-qO-', url), 'xvzC', 'cookbooks/' + name,
                       '--strip-components=1')
            for res in 'cookbooks', 'data_bags', 'environments', 'roles':
                sh.tar('czf', res+'.tar.gz', res)

    def test_chef_solo(self):
        agent_key_file = get_agent_key_file(self.env)
        blueprint_dir = self.blueprint_dir
        self.blueprint_yaml = blueprint_dir / 'chef-solo-test.yaml'
        with YamlPatcher(self.blueprint_yaml) as blueprint:
            bp_info = update_blueprint(self.env, blueprint, 'chef-solo')

        id_ = self.test_id + '-chef-solo-' + str(int(time.time()))
        before, after = self.upload_deploy_and_execute_install(id_, id_)

        fip_node = find_node_state('ip', after['node_state'][id_])
        chef_solo_ip = fip_node['runtimeInfo']['floating_ip_address']

        fabric_env = fabric.api.env
        fabric_env.update({
            'timeout': 30,
            'user': bp_info['users'][0],
            'key_filename': str(agent_key_file),
            'host_string': chef_solo_ip,
        })

        expected_files_contents = (
            ('/tmp/blueprint.txt', 'Great success number #2 !'),
            ('/tmp/blueprint2.txt', '/tmp/blueprint.txt'),
            ('/tmp/chef_node_env.e1.txt', 'env:e1'),
            ('/tmp/chef_node_data_bag_user.db1.i1.txt', 'db1-i1-k1'),
        )

        for file_name, expected_content in expected_files_contents:
            actual_content = fabric.api.run('cat ' + file_name)
            msg = "File '{0}' should have content '{1}' but has '{2}'".format(
                file_name, expected_content, actual_content)
            self.assertEquals(actual_content, expected_content, msg)
<commit_msg>CFY-494 uninstall applications in test and test teardown<commit_after>########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
#    * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    * See the License for the specific language governing permissions and
#    * limitations under the License.

""" Assumes fabric environment already set up """

__author__ = 'ilyash'

import subprocess
import sys
import time
import os

from fabric import operations
import fabric.api
from path import path
import sh

from cosmo_tester.framework.testenv import TestCase
from cosmo_tester.framework.util import YamlPatcher

CHEF_SERVER_COOKBOOK_ZIP_URL = (
    'https://github.com/opscode-cookbooks/chef-server/archive/'
    'c588a4c401d3fac14f70d3285fe49eb4dccd9759.zip'
)

CHEF_SERVER_COOKBOOKS_TAR_GZS = (
    (
        'create-file',
        'https://github.com/ilyash/cookbook-create-file/archive/'
        'bd8218a6a9b3e33a165042241c50908ccb6145d1.tar.gz'
    ),
)

KNIFE_PARAMS = '-u admin -k ~/admin.pem'

IMAGE_NAME = 'Ubuntu 12.04'
FLAVOR_NAME = 'm1.small'

def _use_cookbook(cookbook_name, cookbook_url):
    """ Downloads cookbook from given url and uploads it to the Chef server """
    fabric.api.run('mkdir -p ~/cookbooks/{0}'.format(cookbook_name))
    # Next line was inspired by (sorry, flake8)
    # https://github.com/opscode-cookbooks/chef-server/blame/
    # c588a4c401d3fac14f70d3285fe49eb4dccd9759/README.md#L158
    fabric.api.run('wget -qO- {0} | tar xvzC ~/cookbooks/{1}'
                   ' --strip-components=1'.format(cookbook_url, cookbook_name))
    fabric.api.run('knife cookbook upload {0} --cookbook-path ~/cookbooks {1}'
                   .format(KNIFE_PARAMS, cookbook_name))
    fabric.api.run('knife cookbook list {0} | grep -F {1}'
                   .format(KNIFE_PARAMS, cookbook_name))

def _userize_file(original_path):
    """ Places the file under user's home directory and make it
        permissions-wise accessible """
    fabric.api.sudo("cp -a {path} ~{user}/ && chown {user} ~{user}/{basename}"
                    .format(path=original_path,
                            basename=str(path(original_path).basename()),
                            user=fabric.api.env['user']))

def setup_chef_server(local_dir, cookbooks):
    _userize_file("/etc/chef-server/admin.pem")
    for cb in cookbooks:
        _use_cookbook(*cb)

    _userize_file("/etc/chef-server/chef-validator.pem")
    operations.get('~/chef-validator.pem', str(local_dir))

def find_node_state(node_name, nodes_state):
    pfx = node_name + '_'
    matches = [v for k, v in nodes_state.items() if k.startswith(pfx)]
    if len(matches) != 1:
        raise RuntimeError("Failed to find node {0}".format(node_name))
    return matches[0]

def get_nodes_of_type(blueprint, type_):
    return [n for n in blueprint.obj['blueprint']['nodes']
            if n['type'] == type_]

def get_agent_key_file(env):
    agent_key_file = path(os.path.expanduser(env.agent_key_path)).abspath()
    if not agent_key_file.exists():
        raise RuntimeError("Agent key file {0} does not exist".format(
            agent_key_file))
    return agent_key_file

def update_blueprint(env, blueprint, hostname, userdata_vars=None):
    hostname_base = 'system-test-{0}-{1}'.format(
        time.strftime("%Y%m%d-%H%M"), hostname)
    agent_key_file = get_agent_key_file(env)
    vms = get_nodes_of_type(blueprint, 'cloudify.openstack.server')
    if len(vms) > 1:
        hostnames = ['{0}-{1:2}'.format(hostname_base, i)
                     for i in range(0, len(vms))]
    else:
        hostnames = [hostname_base]

    users = []
    for vm_idx, vm in enumerate(vms):
        vm_hostname = hostnames[vm_idx]
        vm['properties']['worker_config']['key'] = (
            '~/.ssh/' + str(agent_key_file.basename()))

        # vm['properties']['server'] does not exist when using existing one
        if 'server' in vm['properties']:
            vm['properties']['server'].update({
                'flavor_name': FLAVOR_NAME,
                'image_name': IMAGE_NAME,
                'key_name': env.agent_keypair_name,
                'name': vm_hostname,
            })
            vm['properties']['management_network_name'] = (
                env.management_network_name)
            vm['properties']['server']['security_groups'].append(
                env.agents_security_group)
            props = vm['properties']['server']
            if 'userdata' in props:
                props['userdata'] = props['userdata'].format(
                    hostname=vm_hostname, **(userdata_vars or {}))
        users.append(vm['properties']['worker_config']['user'])

    fips = get_nodes_of_type(blueprint, 'cloudify.openstack.floatingip')
    for fip in fips:
            fip_fip = fip['properties']['floatingip']
            fip_fip['floating_network_name'] = env.external_network_name

    return {'hostnames': hostnames, 'users': users}

class ChefPluginClientTest(TestCase):

    def setUp(self, *args, **kwargs):

        super(ChefPluginClientTest, self).setUp(*args, **kwargs)
        agent_key_file = get_agent_key_file(self.env)

        blueprint_dir = self.copy_blueprint('chef-plugin')
        self.blueprint_yaml = blueprint_dir / 'chef-server-by-chef-solo.yaml'

        with YamlPatcher(self.blueprint_yaml) as blueprint:
            bp_info = update_blueprint(self.env, blueprint, 'chef-server')

        self.chef_server_hostname = bp_info['hostnames'][0]

        cookbooks_dir = blueprint_dir / 'cookbooks'
        cookbooks_dir.mkdir()

        def run(*args, **kwargs):
            return subprocess.check_output(*args, stderr=sys.stderr, **kwargs)

        with cookbooks_dir:
            run([
                'wget', '-q', '-O', 'chef-server.zip',
                CHEF_SERVER_COOKBOOK_ZIP_URL,
                ])
            run(['unzip', 'chef-server.zip'])
            chef_cookbook_dir = cookbooks_dir.glob('chef-server-*')[0]
            run(['mv', chef_cookbook_dir, 'chef-server'])
            # Next line because Chef cookbooks are required
            # to declare all dependencies, even if they don't use them.
            # We don't need git, it's only used in chef-cookbook::dev recipe.
            run(['sed', '-i', "/depends 'git'/d", 'chef-server/metadata.rb'])

        with blueprint_dir:
            run(['tar', 'czf', 'cookbooks.tar.gz', 'cookbooks'])

        self.chef_server_id = self.test_id + '-chef-server'
        id_ = self.chef_server_id
        before, after = self.upload_deploy_and_execute_install(id_, id_)

        fip_node = find_node_state('ip', after['node_state'][id_])
        self.chef_server_ip = fip_node['runtimeInfo']['floating_ip_address']

        fabric_env = fabric.api.env
        fabric_env.update({
            'timeout': 30,
            'user': bp_info['users'][0],
            'key_filename': str(agent_key_file),
            'host_string': self.chef_server_ip,
        })

        setup_chef_server(blueprint_dir, CHEF_SERVER_COOKBOOKS_TAR_GZS)
        self.blueprint_dir = blueprint_dir

    def tearDown(self, *args, **kwargs):
        self.execute_uninstall(self.chef_server_id)
        super(ChefPluginClientTest, self).tearDown(*args, **kwargs)

    def test_chef_client(self):
        blueprint_dir = self.blueprint_dir
        self.blueprint_yaml = blueprint_dir / 'chef-client-test.yaml'
        with YamlPatcher(self.blueprint_yaml) as blueprint:
            update_blueprint(self.env, blueprint, 'chef-server', {
                'chef_server_ip': self.chef_server_ip,
                'chef_server_hostname': self.chef_server_hostname,
            })
            chef_node = get_nodes_of_type(blueprint,
                                          'cloudify.types.chef.db_server')[0]
            chef_config = chef_node['properties']['chef_config']
            chef_config['chef_server_url'] = 'https://{0}:443'.format(
                self.chef_server_ip)
            chef_config['validation_client_name'] = 'chef-validator'
            chef_config['validation_key'] = (
                path(blueprint_dir) / 'chef-validator.pem').text()

        id_ = self.test_id + '-chef-client-' + str(int(time.time()))
        before, after = self.upload_deploy_and_execute_install(id_, id_)

        fip_node = find_node_state('ip', after['node_state'][id_])
        chef_client_ip = fip_node['runtimeInfo']['floating_ip_address']

        fabric_env = fabric.api.env
        fabric_env.update({
            # XXX: sometime - same user for connection is accidental
            #      todo: replace it with update_blueprint()'s bp_info,
            #            as in setUp()
            'host_string': chef_client_ip,
        })

        out = fabric.api.run('cat /tmp/blueprint.txt')
        self.assertEquals(out, 'Great success!')

        self.execute_uninstall(id_)

class ChefPluginSoloTest(TestCase):

    def setUp(self, *args, **kwargs):

        super(ChefPluginSoloTest, self).setUp(*args, **kwargs)

        self.blueprint_dir = self.copy_blueprint('chef-plugin')

        # Get resources
        with self.blueprint_dir:
            # Cookbooks
            for name, url in CHEF_SERVER_COOKBOOKS_TAR_GZS:
                sh.mkdir('-p', 'cookbooks/' + name)
                sh.tar(sh.wget('-qO-', url), 'xvzC', 'cookbooks/' + name,
                       '--strip-components=1')
            for res in 'cookbooks', 'data_bags', 'environments', 'roles':
                sh.tar('czf', res+'.tar.gz', res)

    def test_chef_solo(self):
        agent_key_file = get_agent_key_file(self.env)
        blueprint_dir = self.blueprint_dir
        self.blueprint_yaml = blueprint_dir / 'chef-solo-test.yaml'
        with YamlPatcher(self.blueprint_yaml) as blueprint:
            bp_info = update_blueprint(self.env, blueprint, 'chef-solo')

        id_ = self.test_id + '-chef-solo-' + str(int(time.time()))
        before, after = self.upload_deploy_and_execute_install(id_, id_)

        fip_node = find_node_state('ip', after['node_state'][id_])
        chef_solo_ip = fip_node['runtimeInfo']['floating_ip_address']

        fabric_env = fabric.api.env
        fabric_env.update({
            'timeout': 30,
            'user': bp_info['users'][0],
            'key_filename': str(agent_key_file),
            'host_string': chef_solo_ip,
        })

        expected_files_contents = (
            ('/tmp/blueprint.txt', 'Great success number #2 !'),
            ('/tmp/blueprint2.txt', '/tmp/blueprint.txt'),
            ('/tmp/chef_node_env.e1.txt', 'env:e1'),
            ('/tmp/chef_node_data_bag_user.db1.i1.txt', 'db1-i1-k1'),
        )

        for file_name, expected_content in expected_files_contents:
            actual_content = fabric.api.run('cat ' + file_name)
            msg = "File '{0}' should have content '{1}' but has '{2}'".format(
                file_name, expected_content, actual_content)
            self.assertEquals(actual_content, expected_content, msg)

        self.execute_uninstall(id_)<|endoftext|>


<commit_before>from __future__ import division

__author__ = "John Kirkham <kirkhamj@janelia.hhmi.org>"
__date__ = "$Aug 09, 2017 17:12$"

import itertools

import numpy

import dask
import dask.array
import dask.array.linalg

from builtins import range as irange

from nanshe.imp.filters.wavelet import transform as _wavelet_transform
from nanshe.imp.segment import extract_f0 as _extract_f0

def extract_f0(new_data,
               half_window_size,
               which_quantile,
               temporal_smoothing_gaussian_filter_stdev,
               temporal_smoothing_gaussian_filter_window_size,
               spatial_smoothing_gaussian_filter_stdev,
               spatial_smoothing_gaussian_filter_window_size,
               bias=None,
               return_f0=False,
               **parameters):
    """
        Compute ``extract_f0`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    if "out" in parameters:
        raise TypeError("Got an unexpected keyword argument 'out'.")

    new_data = new_data.astype(numpy.float32)

    if bias is None:
        bias = 1 - new_data.min()

    depth = numpy.zeros((new_data.ndim,), dtype=int)

    depth[0] = int(numpy.ceil(
        temporal_smoothing_gaussian_filter_window_size *
        temporal_smoothing_gaussian_filter_stdev
    ))
    depth[1:] = int(numpy.ceil(
        spatial_smoothing_gaussian_filter_window_size *
        spatial_smoothing_gaussian_filter_stdev
    ))
    depth[0] = max(depth[0], half_window_size)

    depth = tuple(depth.tolist())

    boundary = len(depth) * ["none"]
    # Workaround for a bug in Dask with 0 depth.
    #
    # ref: https://github.com/dask/dask/issues/2258
    #
    for i in irange(len(depth)):
        if boundary[i] == "none" and depth[i] == 0:
            boundary[i] = "reflect"

    boundary = tuple(boundary)

    result = new_data.map_overlap(
        func=_extract_f0,
        depth=depth,
        boundary=boundary,
        dtype=numpy.float32,
        half_window_size=half_window_size,
        which_quantile=which_quantile,
        temporal_smoothing_gaussian_filter_stdev=temporal_smoothing_gaussian_filter_stdev,
        temporal_smoothing_gaussian_filter_window_size=temporal_smoothing_gaussian_filter_window_size,
        spatial_smoothing_gaussian_filter_stdev=spatial_smoothing_gaussian_filter_stdev,
        spatial_smoothing_gaussian_filter_window_size=spatial_smoothing_gaussian_filter_window_size,
        bias=bias,
        return_f0=return_f0,
        **parameters
    )

    return(result)

def wavelet_transform(im0,
                      scale=5):
    """
        Compute ``wavelet_transform`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    im0 = im0.astype(numpy.float32)

    try:
        scale_iter = enumerate(scale)
    except TypeError:
        scale_iter = enumerate(itertools.repeat(scale, im0.ndim))

    depth = list(itertools.repeat(0, im0.ndim))
    for i, each_scale in scale_iter:
        depth_i = 0
        for j in irange(1, 1 + each_scale):
            depth_i += 2 ** j

        depth[i] = depth_i

    depth = tuple(depth)

    boundary = len(depth) * ["none"]
    # Workaround for a bug in Dask with 0 depth.
    #
    # ref: https://github.com/dask/dask/issues/2258
    #
    for i in irange(len(depth)):
        if boundary[i] == "none" and depth[i] == 0:
            boundary[i] = "reflect"

    boundary = tuple(boundary)

    result = im0.map_overlap(
        func=_wavelet_transform,
        depth=depth,
        boundary=boundary,
        dtype=numpy.float32,
        scale=scale
    )

    return(result)

def zeroed_mean_images(new_data):
    """
        Compute ``zeroed_mean_images`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    new_data_means = new_data.mean(
        axis=tuple(irange(1, new_data.ndim)),
        keepdims=True
    )

    new_data_zeroed = new_data - new_data_means

    return new_data_zeroed

def renormalized_images(input_array, ord=2):
    """
        Compute ``renormalized_images`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    input_array_norms = dask.array.linalg.norm(
        input_array.reshape((len(input_array), -1)),
        ord=ord,
        axis=1
    )
    input_array_norms = input_array_norms[
        (slice(None),) + (input_array.ndim - input_array_norms.ndim) * (None,)
    ]

    input_array_renormed = dask.array.where(
        input_array_norms != 0, input_array / input_array_norms, input_array
    )

    return(input_array_renormed)

def normalize_data(new_data, **parameters):
    """
        Compute ``normalize_data`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    if "out" in parameters:
        raise TypeError("Got an unexpected keyword argument 'out'.")

    ord = parameters.get("renormalized_images", {}).get("ord", 2)

    if ord is None:
        ord = 2

    new_data_zeroed = zeroed_mean_images(new_data)

    new_data_norms = dask.array.linalg.norm(
        new_data_zeroed.reshape((len(new_data_zeroed), -1)),
        ord=ord,
        axis=1
    )
    new_data_norms = new_data_norms[
        (slice(None),) + (new_data_zeroed.ndim - new_data_norms.ndim) * (None,)
    ]

    new_data_renormed = dask.array.where(
        new_data_norms != 0, new_data_zeroed / new_data_norms, new_data_zeroed
    )

    return(new_data_renormed)
<commit_msg>Use renormalized_images in normalize_data<commit_after>from __future__ import division

__author__ = "John Kirkham <kirkhamj@janelia.hhmi.org>"
__date__ = "$Aug 09, 2017 17:12$"

import itertools

import numpy

import dask
import dask.array
import dask.array.linalg

from builtins import range as irange

from nanshe.imp.filters.wavelet import transform as _wavelet_transform
from nanshe.imp.segment import extract_f0 as _extract_f0

def extract_f0(new_data,
               half_window_size,
               which_quantile,
               temporal_smoothing_gaussian_filter_stdev,
               temporal_smoothing_gaussian_filter_window_size,
               spatial_smoothing_gaussian_filter_stdev,
               spatial_smoothing_gaussian_filter_window_size,
               bias=None,
               return_f0=False,
               **parameters):
    """
        Compute ``extract_f0`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    if "out" in parameters:
        raise TypeError("Got an unexpected keyword argument 'out'.")

    new_data = new_data.astype(numpy.float32)

    if bias is None:
        bias = 1 - new_data.min()

    depth = numpy.zeros((new_data.ndim,), dtype=int)

    depth[0] = int(numpy.ceil(
        temporal_smoothing_gaussian_filter_window_size *
        temporal_smoothing_gaussian_filter_stdev
    ))
    depth[1:] = int(numpy.ceil(
        spatial_smoothing_gaussian_filter_window_size *
        spatial_smoothing_gaussian_filter_stdev
    ))
    depth[0] = max(depth[0], half_window_size)

    depth = tuple(depth.tolist())

    boundary = len(depth) * ["none"]
    # Workaround for a bug in Dask with 0 depth.
    #
    # ref: https://github.com/dask/dask/issues/2258
    #
    for i in irange(len(depth)):
        if boundary[i] == "none" and depth[i] == 0:
            boundary[i] = "reflect"

    boundary = tuple(boundary)

    result = new_data.map_overlap(
        func=_extract_f0,
        depth=depth,
        boundary=boundary,
        dtype=numpy.float32,
        half_window_size=half_window_size,
        which_quantile=which_quantile,
        temporal_smoothing_gaussian_filter_stdev=temporal_smoothing_gaussian_filter_stdev,
        temporal_smoothing_gaussian_filter_window_size=temporal_smoothing_gaussian_filter_window_size,
        spatial_smoothing_gaussian_filter_stdev=spatial_smoothing_gaussian_filter_stdev,
        spatial_smoothing_gaussian_filter_window_size=spatial_smoothing_gaussian_filter_window_size,
        bias=bias,
        return_f0=return_f0,
        **parameters
    )

    return(result)

def wavelet_transform(im0,
                      scale=5):
    """
        Compute ``wavelet_transform`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    im0 = im0.astype(numpy.float32)

    try:
        scale_iter = enumerate(scale)
    except TypeError:
        scale_iter = enumerate(itertools.repeat(scale, im0.ndim))

    depth = list(itertools.repeat(0, im0.ndim))
    for i, each_scale in scale_iter:
        depth_i = 0
        for j in irange(1, 1 + each_scale):
            depth_i += 2 ** j

        depth[i] = depth_i

    depth = tuple(depth)

    boundary = len(depth) * ["none"]
    # Workaround for a bug in Dask with 0 depth.
    #
    # ref: https://github.com/dask/dask/issues/2258
    #
    for i in irange(len(depth)):
        if boundary[i] == "none" and depth[i] == 0:
            boundary[i] = "reflect"

    boundary = tuple(boundary)

    result = im0.map_overlap(
        func=_wavelet_transform,
        depth=depth,
        boundary=boundary,
        dtype=numpy.float32,
        scale=scale
    )

    return(result)

def zeroed_mean_images(new_data):
    """
        Compute ``zeroed_mean_images`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    new_data_means = new_data.mean(
        axis=tuple(irange(1, new_data.ndim)),
        keepdims=True
    )

    new_data_zeroed = new_data - new_data_means

    return new_data_zeroed

def renormalized_images(input_array, ord=2):
    """
        Compute ``renormalized_images`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    input_array_norms = dask.array.linalg.norm(
        input_array.reshape((len(input_array), -1)),
        ord=ord,
        axis=1
    )
    input_array_norms = input_array_norms[
        (slice(None),) + (input_array.ndim - input_array_norms.ndim) * (None,)
    ]

    input_array_renormed = dask.array.where(
        input_array_norms != 0, input_array / input_array_norms, input_array
    )

    return(input_array_renormed)

def normalize_data(new_data, **parameters):
    """
        Compute ``normalize_data`` on Dask Arrays.

        See the nanshe function for more details

        Returns:
            Dask Array:    A lazily computed result.
    """

    if "out" in parameters:
        raise TypeError("Got an unexpected keyword argument 'out'.")

    new_data_zeroed = zeroed_mean_images(new_data)

    new_data_renormed = renormalized_images(
        new_data_zeroed, **parameters.get("renormalized_images", {})
    )

    return(new_data_renormed)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# Copyright (C) 2019 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os

from sys import exit
from copy import deepcopy
from netifaces import interfaces

from vyos.config import Config
from vyos.ifconfig import GeneveIf
from vyos.validate import is_member
from vyos import ConfigError

default_config_data = {
    'address': [],
    'deleted': False,
    'description': '',
    'disable': False,
    'intf': '',
    'ip_arp_cache_tmo': 30,
    'ip_proxy_arp': 0,
    'is_bridge_member': False,
    'mtu': 1500,
    'remote': '',
    'vni': ''
}

def get_config():
    geneve = deepcopy(default_config_data)
    conf = Config()

    # determine tagNode instance
    if 'VYOS_TAGNODE_VALUE' not in os.environ:
        raise ConfigError('Interface (VYOS_TAGNODE_VALUE) not specified')

    geneve['intf'] = os.environ['VYOS_TAGNODE_VALUE']

    # check if interface is member if a bridge
    geneve['is_bridge_member'] = is_member(conf, geneve['intf'], 'bridge')

    # Check if interface has been removed
    if not conf.exists('interfaces geneve ' + geneve['intf']):
        geneve['deleted'] = True
        return geneve

    # set new configuration level
    conf.set_level('interfaces geneve ' + geneve['intf'])

    # retrieve configured interface addresses
    if conf.exists('address'):
        geneve['address'] = conf.return_values('address')

    # retrieve interface description
    if conf.exists('description'):
        geneve['description'] = conf.return_value('description')

    # Disable this interface
    if conf.exists('disable'):
        geneve['disable'] = True

    # ARP cache entry timeout in seconds
    if conf.exists('ip arp-cache-timeout'):
        geneve['ip_arp_cache_tmo'] = int(conf.return_value('ip arp-cache-timeout'))

    # Enable proxy-arp on this interface
    if conf.exists('ip enable-proxy-arp'):
        geneve['ip_proxy_arp'] = 1

    # Maximum Transmission Unit (MTU)
    if conf.exists('mtu'):
        geneve['mtu'] = int(conf.return_value('mtu'))

    # Remote address of GENEVE tunnel
    if conf.exists('remote'):
        geneve['remote'] = conf.return_value('remote')

    # Virtual Network Identifier
    if conf.exists('vni'):
        geneve['vni'] = conf.return_value('vni')

    return geneve

def verify(geneve):
    if geneve['deleted']:
        if geneve['is_bridge_member']:
            interface = geneve['intf']
            bridge = geneve['is_bridge_member']
            raise ConfigError(f'Interface "{interface}" can not be deleted as it belongs to bridge "{bridge}"!')

        return None

    if geneve['is_bridge_member'] and geneve['address']:
        raise ConfigError((
            f'Cannot assign address to interface "{geneve["intf"]}" '
            f'as it is a member of bridge "{geneve["is_bridge_member"]}"!'))

    if not geneve['remote']:
        raise ConfigError('GENEVE remote must be configured')

    if not geneve['vni']:
        raise ConfigError('GENEVE VNI must be configured')

    return None

def generate(geneve):
    return None

def apply(geneve):
    # Check if GENEVE interface already exists
    if geneve['intf'] in interfaces():
        g = GeneveIf(geneve['intf'])
        # GENEVE is super picky and the tunnel always needs to be recreated,
        # thus we can simply always delete it first.
        g.remove()

    if not geneve['deleted']:
        # GENEVE interface needs to be created on-block
        # instead of passing a ton of arguments, I just use a dict
        # that is managed by vyos.ifconfig
        conf = deepcopy(GeneveIf.get_config())

        # Assign GENEVE instance configuration parameters to config dict
        conf['vni'] = geneve['vni']
        conf['remote'] = geneve['remote']

        # Finally create the new interface
        g = GeneveIf(geneve['intf'], **conf)
        # update interface description used e.g. by SNMP
        g.set_alias(geneve['description'])
        # Maximum Transfer Unit (MTU)
        g.set_mtu(geneve['mtu'])

        # configure ARP cache timeout in milliseconds
        g.set_arp_cache_tmo(geneve['ip_arp_cache_tmo'])
        # Enable proxy-arp on this interface
        g.set_proxy_arp(geneve['ip_proxy_arp'])

        # Configure interface address(es) - no need to implicitly delete the
        # old addresses as they have already been removed by deleting the
        # interface above
        for addr in geneve['address']:
            g.add_addr(addr)

        # As the GENEVE interface is always disabled first when changing
        # parameters we will only re-enable the interface if it is not
        # administratively disabled
        if not geneve['disable']:
            g.set_admin_state('up')

        # re-add ourselves to any bridge we might have fallen out of
        if geneve['is_bridge_member']:
            g.add_to_bridge(geneve['is_bridge_member'])

    return None

if __name__ == '__main__':
    try:
        c = get_config()
        verify(c)
        generate(c)
        apply(c)
    except ConfigError as e:
        print(e)
        exit(1)
<commit_msg>geneve: T2241: cleanup verify section<commit_after>#!/usr/bin/env python3
#
# Copyright (C) 2019 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os

from sys import exit
from copy import deepcopy
from netifaces import interfaces

from vyos.config import Config
from vyos.ifconfig import GeneveIf
from vyos.validate import is_member
from vyos import ConfigError

default_config_data = {
    'address': [],
    'deleted': False,
    'description': '',
    'disable': False,
    'intf': '',
    'ip_arp_cache_tmo': 30,
    'ip_proxy_arp': 0,
    'is_bridge_member': False,
    'mtu': 1500,
    'remote': '',
    'vni': ''
}

def get_config():
    geneve = deepcopy(default_config_data)
    conf = Config()

    # determine tagNode instance
    if 'VYOS_TAGNODE_VALUE' not in os.environ:
        raise ConfigError('Interface (VYOS_TAGNODE_VALUE) not specified')

    geneve['intf'] = os.environ['VYOS_TAGNODE_VALUE']

    # check if interface is member if a bridge
    geneve['is_bridge_member'] = is_member(conf, geneve['intf'], 'bridge')

    # Check if interface has been removed
    if not conf.exists('interfaces geneve ' + geneve['intf']):
        geneve['deleted'] = True
        return geneve

    # set new configuration level
    conf.set_level('interfaces geneve ' + geneve['intf'])

    # retrieve configured interface addresses
    if conf.exists('address'):
        geneve['address'] = conf.return_values('address')

    # retrieve interface description
    if conf.exists('description'):
        geneve['description'] = conf.return_value('description')

    # Disable this interface
    if conf.exists('disable'):
        geneve['disable'] = True

    # ARP cache entry timeout in seconds
    if conf.exists('ip arp-cache-timeout'):
        geneve['ip_arp_cache_tmo'] = int(conf.return_value('ip arp-cache-timeout'))

    # Enable proxy-arp on this interface
    if conf.exists('ip enable-proxy-arp'):
        geneve['ip_proxy_arp'] = 1

    # Maximum Transmission Unit (MTU)
    if conf.exists('mtu'):
        geneve['mtu'] = int(conf.return_value('mtu'))

    # Remote address of GENEVE tunnel
    if conf.exists('remote'):
        geneve['remote'] = conf.return_value('remote')

    # Virtual Network Identifier
    if conf.exists('vni'):
        geneve['vni'] = conf.return_value('vni')

    return geneve

def verify(geneve):
    if geneve['deleted']:
        if geneve['is_bridge_member']:
            raise ConfigError((
                f'Cannot delete interface "{geneve["intf"]}" as it is a '
                f'member of bridge "{geneve["is_bridge_member"]}"!'))

        return None

    if geneve['is_bridge_member'] and geneve['address']:
        raise ConfigError((
            f'Cannot assign address to interface "{geneve["intf"]}" '
            f'as it is a member of bridge "{geneve["is_bridge_member"]}"!'))

    if not geneve['remote']:
        raise ConfigError('GENEVE remote must be configured')

    if not geneve['vni']:
        raise ConfigError('GENEVE VNI must be configured')

    return None

def generate(geneve):
    return None

def apply(geneve):
    # Check if GENEVE interface already exists
    if geneve['intf'] in interfaces():
        g = GeneveIf(geneve['intf'])
        # GENEVE is super picky and the tunnel always needs to be recreated,
        # thus we can simply always delete it first.
        g.remove()

    if not geneve['deleted']:
        # GENEVE interface needs to be created on-block
        # instead of passing a ton of arguments, I just use a dict
        # that is managed by vyos.ifconfig
        conf = deepcopy(GeneveIf.get_config())

        # Assign GENEVE instance configuration parameters to config dict
        conf['vni'] = geneve['vni']
        conf['remote'] = geneve['remote']

        # Finally create the new interface
        g = GeneveIf(geneve['intf'], **conf)
        # update interface description used e.g. by SNMP
        g.set_alias(geneve['description'])
        # Maximum Transfer Unit (MTU)
        g.set_mtu(geneve['mtu'])

        # configure ARP cache timeout in milliseconds
        g.set_arp_cache_tmo(geneve['ip_arp_cache_tmo'])
        # Enable proxy-arp on this interface
        g.set_proxy_arp(geneve['ip_proxy_arp'])

        # Configure interface address(es) - no need to implicitly delete the
        # old addresses as they have already been removed by deleting the
        # interface above
        for addr in geneve['address']:
            g.add_addr(addr)

        # As the GENEVE interface is always disabled first when changing
        # parameters we will only re-enable the interface if it is not
        # administratively disabled
        if not geneve['disable']:
            g.set_admin_state('up')

        # re-add ourselves to any bridge we might have fallen out of
        if geneve['is_bridge_member']:
            g.add_to_bridge(geneve['is_bridge_member'])

    return None

if __name__ == '__main__':
    try:
        c = get_config()
        verify(c)
        generate(c)
        apply(c)
    except ConfigError as e:
        print(e)
        exit(1)
<|endoftext|>


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

import argparse
import datetime
import errno
import os
import sys
import time

from splinter import Browser

class Color:
    # Use colors to make command line output prettier
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    END = '\033[0m'

class Shuriken:

    def __init__(self):
        # Expect some weird characters from fuzz lists, make encoding UTF-8
        reload(sys)
        sys.setdefaultencoding('utf8')

        # All potential XSS findings
        self.xss_links = []

        # Get user args and store
        self.user_args = self.parse_args()

        # PhantomJS browser
        self.browser = Browser("phantomjs")

    def main(self):
        # Print out a welcome message
        print "\n"
        print "=" * 34 + Color.YELLOW + "\nWelcome to the" + Color.RED + \
            " Shuriken " + Color.YELLOW + "XSS tool!\n" + \
            Color.END + "=" * 34 + "\n"

        self.test_xss(self.user_args.PAYLOADS_LIST,
                      self.user_args.URL,
                      self.user_args.REQUEST_DELAY,
                      self.user_args.SCREENSHOT_NAME)
        print Color.GREEN + "\n=== Testing complete! ===\n" + Color.END

        # If the test found possible XSS vulnerabilities, ask if we should
        # log
        if self.xss_links:
            print Color.YELLOW + \
                "Potential XSS vulnerabilities were detected!" + \
                Color.END
            self.log_file(self.xss_links)
        else:
            print Color.YELLOW + \
                "No potential XSS vulnerabilities detected...\n" + \
                Color.END
            print "Goodbye!"
            print "\n"

    def make_sure_path_exists(self, path):
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise

    def inject_payload(self, payload, link, request_delay, screenshot_target):
        # Visit user supplied link with injected payload
        browser = self.browser

        # Let user specify where in the URL fuzz values should be injected
        injected_link = link.replace("{xss}", payload)

        # If user added a delay, wait that amount of time before requesting
        # Also, if user wants to interrupt/exit during wait time,
        # let them log before leaving
        if request_delay is not None:
            try:
                time.sleep(float(request_delay))
            except KeyboardInterrupt:
                print Color.YELLOW + \
                    "\nTesting interrupted by user!\n" + Color.END
                self.log_file(self.xss_links)
                sys.exit()

        browser.visit(injected_link)

        # Keep index of screens, so they can be easily
        # linked to line nums in log
        self.screen_index = str(len(self.xss_links) + 1)

        # Check to see if payload was reflected in HTML source
        if payload in browser.html:
            print Color.GREEN + "\n[+] Potential XSS vulnerability found:" + \
                Color.END
            # If user set the --screen flag to target, capture screen of
            # payload
            if screenshot_target is not None:
                # Check if screenshots directory exists, if not then create it
                self.make_sure_path_exists("screenshots")
                screenshot_file_name = "screenshots/" + \
                    screenshot_target + "_" + \
                    datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + \
                    "_" + self.screen_index + ".png"
                # Save screenshot to directory
                browser.driver.save_screenshot(screenshot_file_name)
                print Color.YELLOW + "Screenshot saved: " + \
                    screenshot_file_name + Color.END
            # Add link to list of all positive XSS hits
            self.xss_links.append(injected_link)
            return Color.BLUE + injected_link + Color.END
        else:
            return Color.YELLOW + "\n[+] Tested, but no XSS found at: \n" + \
                Color.RED + injected_link + Color.END

    def test_xss(self, payloads_param, link, request_delay, screenshot_target):
        # If the user added time delay, show them what it is set to.
        if request_delay is not None:
            print Color.YELLOW + "\n[!] Request delay is set to [" + \
                str(request_delay) + "] seconds between requests." + Color.END

        # Load the payload file and inject all payloads
        # into user supplied URL to test for XSS
        payloads = []
        with open(payloads_param) as file:
            for line in file:
                line = line.strip()
                payloads.append(line)
        for item in payloads:
            print self.inject_payload(item, link, request_delay,
                                      screenshot_target)

    def log_file(self, link_list):
        # Prompt the user to confirm log file, if yes, log XSS hits
        log_confirm = raw_input(
            "\nWould you like to save these results? [y/n] > ")
        if log_confirm == "y":
            target_name = raw_input("Please enter the target name > ")
            # Check if logs directory exists, if not then create it
            self.make_sure_path_exists("logs")
            # Save log file to directory
            file_name = "logs/" + target_name + "_" + \
                datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + \
                ".txt"
            with open(file_name, 'w') as file:
                for link in link_list:
                    file.write(link)
                    file.write("\n")
                # Add metadata about what payload file was used
                file.write("\n*** Created from the payload file >>> " +
                           self.user_args.PAYLOADS_LIST)
                file.close()
            print "\nFile successfully saved as: " + \
                Color.BLUE + file_name + Color.END
            print "\n"
        else:
            print "\nGoodbye!"
            print "\n"

    def parse_args(self):
        # Parse arguments from the user sent on command line
        parser = argparse.ArgumentParser()
        parser.add_argument(
            '-u', action='store', dest='URL',
            help='The URL to analyze', required=True)
        parser.add_argument(
            '-p', action='store', dest='PAYLOADS_LIST',
            help='The payload list to use', required=True)
        parser.add_argument(
            '-t', action='store', dest='REQUEST_DELAY',
            help='Amount of time in seconds to delay between requests')
        parser.add_argument(
            '--screen', action='store', dest='SCREENSHOT_NAME',
            help='Screens of target')

        arguments = parser.parse_args()

        # Check for existence of '{xss}' injection point in URL string
        if "{xss}" not in arguments.URL:
            print Color.RED + "Please provide the '{xss}' placeholder for" + \
                " injection point in the URL" + Color.END
            print Color.GREEN + \
                "Example: -u \"http://example.com/index.php?name={xss}\"" + \
                Color.END
            exit()

        return arguments

if __name__ == "__main__":
    try:
        Shuriken().main()
    except KeyboardInterrupt:
        print Color.YELLOW + \
            "\nTesting interrupted by user!\n" + Color.END
        Shuriken().log_file(Shuriken().xss_links)
        sys.exit()
<commit_msg>Restructure to be more modular<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import datetime
import errno
import os
import sys
import time

from splinter import Browser

class Color:
    # Use colors to make command line output prettier
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    END = '\033[0m'

class Shuriken:

    def __init__(self):
        # Expect some weird characters from fuzz lists, make encoding UTF-8
        reload(sys)
        sys.setdefaultencoding('utf8')

        # All potential XSS findings
        self.xss_links = []

        # Get user args and store
        self.user_args = self.parse_args()

        # Keep index of screens, so they can be easily
        # linked to line nums in log
        self.screen_index = str(len(self.xss_links) + 1)

        # PhantomJS browser
        self.browser = Browser("phantomjs")

    def main(self):
        # Print out a welcome message
        print "\n"
        print "=" * 34 + Color.YELLOW + "\nWelcome to the" + Color.RED + \
            " Shuriken " + Color.YELLOW + "XSS tool!\n" + \
            Color.END + "=" * 34 + "\n"

        self.test_xss(self.user_args.PAYLOADS_LIST,
                      self.user_args.URL,
                      self.user_args.REQUEST_DELAY,
                      self.user_args.SCREENSHOT_NAME)
        print Color.GREEN + "\n=== Testing complete! ===\n" + Color.END

        # If the test found possible XSS vulnerabilities, ask if we should
        # log
        if self.xss_links:
            print Color.YELLOW + \
                "Potential XSS vulnerabilities were detected!" + \
                Color.END
            self.log_file(self.xss_links)
        else:
            print Color.YELLOW + \
                "No potential XSS vulnerabilities detected...\n" + \
                Color.END
            print "Goodbye!"
            print "\n"

    def make_sure_path_exists(self, path):
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise

    def inject_payload(self, payload, link, request_delay, screenshot_target):
        # Visit user supplied link with injected payload
        browser = self.browser

        # Let user specify where in the URL fuzz values should be injected
        injected_link = link.replace("{xss}", payload)

        # If user added a delay, wait that amount of time before requesting
        # Also, if user wants to interrupt/exit during wait time,
        # let them log before leaving
        if request_delay is not None:
            try:
                time.sleep(float(request_delay))
            except KeyboardInterrupt:
                print Color.YELLOW + \
                    "\nTesting interrupted by user!\n" + Color.END
                self.log_file(self.xss_links)
                sys.exit()

        browser.visit(injected_link)

        # Check to see if payload was reflected in HTML source
        print self.detect_xss(payload, browser, screenshot_target, injected_link)

    def detect_xss(self, payload, browser_object, screenshot_target, injected_link):
        # Check to see if payload was reflected in HTML source
        if payload in browser_object.html_source:
            print Color.GREEN + "\n[+] Potential XSS vulnerability found:" + \
                Color.END
            # If user set the --screen flag to target, capture screenshot of
            # payload
            if screenshot_target is not None:
                self.take_screenshot(screenshot_target, browser_object)
            # Add link to list of all positive XSS hits
            self.xss_links.append(injected_link)
            return Color.BLUE + injected_link + Color.END
        else:
            return Color.YELLOW + "\n[+] Tested, but no XSS found at: \n" + \
                Color.RED + injected_link + Color.END

    def take_screenshot(self, screenshot_target, browser_object):
        # Check if screenshots directory exists, if not then create it
        self.make_sure_path_exists("screenshots")
        screenshot_file_name = "screenshots/" + \
            screenshot_target + "_" + \
            datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + \
            "_" + self.screen_index + ".png"
        # Save screenshot to directory
        browser_object.driver.save_screenshot(screenshot_file_name)
        print Color.YELLOW + "Screenshot saved: " + \
            screenshot_file_name + Color.END

    def test_xss(self, payloads_param, link, request_delay, screenshot_target):
        # If the user added time delay, show them what it is set to.
        if request_delay is not None:
            print Color.YELLOW + "\n[!] Request delay is set to [" + \
                str(request_delay) + "] seconds between requests." + Color.END

        # Load the payload file and inject all payloads
        # into user supplied URL to test for XSS
        payloads = []
        with open(payloads_param) as payload_file:
            for line in payload_file:
                line = line.strip()
                payloads.append(line)
        for item in payloads:
            print self.inject_payload(item, link, request_delay,
                                      screenshot_target)

    def log_file(self, link_list):
        # Prompt the user to confirm log file, if yes, log XSS hits
        log_confirm = raw_input(
            "\nWould you like to save these results? [y/n] > ")
        if log_confirm == "y":
            target_name = raw_input("Please enter the target name > ")
            # Check if logs directory exists, if not then create it
            self.make_sure_path_exists("logs")
            # Save log file to directory
            file_name = "logs/" + target_name + "_" + \
                datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + \
                ".txt"
            with open(file_name, 'w') as file:
                for link in link_list:
                    file.write(link)
                    file.write("\n")
                # Add metadata about what payload file was used
                file.write("\n*** Created from the payload file >>> " +
                           self.user_args.PAYLOADS_LIST)
                file.close()
            print "\nFile successfully saved as: " + \
                Color.BLUE + file_name + Color.END
            print "\n"
        else:
            print "\nGoodbye!"
            print "\n"

    def parse_args(self):
        # Parse arguments from the user sent on command line
        parser = argparse.ArgumentParser()
        parser.add_argument(
            '-u', action='store', dest='URL',
            help='The URL to analyze', required=True)
        parser.add_argument(
            '-p', action='store', dest='PAYLOADS_LIST',
            help='The payload list to use', required=True)
        parser.add_argument(
            '-t', action='store', dest='REQUEST_DELAY',
            help='Amount of time in seconds to delay between requests')
        parser.add_argument(
            '--screen', action='store', dest='SCREENSHOT_NAME',
            help='Screens of target')

        arguments = parser.parse_args()

        # Check for existence of '{xss}' injection point in URL string
        if "{xss}" not in arguments.URL:
            print Color.RED + "Please provide the '{xss}' placeholder for" + \
                " injection point in the URL" + Color.END
            print Color.GREEN + \
                "Example: -u \"http://example.com/index.php?name={xss}\"" + \
                Color.END
            exit()

        return arguments

if __name__ == "__main__":
    try:
        Shuriken().main()
    except KeyboardInterrupt:
        print Color.YELLOW + \
            "\nTesting interrupted by user!\n" + Color.END
        Shuriken().log_file(Shuriken().xss_links)
        sys.exit()
<|endoftext|>