<commit_before>from heapq import heappush, heapreplace
from collections import Counter
import array

def chunkpos_iter(nsize: int, lseq: int, w: int) -> (int, int):
    """
    Iterator of chunk indices.

    - nsize: n in ngram
    - lseq: length of sequence to chunk
    - w: width of buffer

    """

    assert nsize <= w
    ew = w-nsize+1

    nchunks = (lseq // ew)

    leftover = lseq - nchunks*ew

    if leftover >= nsize:
        nchunks += 1

    for w_i in range(nchunks):
        slice_beg = (w_i*ew)
        slice_end = slice_beg + w
        yield (slice_beg, slice_end)

class MaxHashNgramSketch(object):

    """
    MaxHash Sketch.
    """

    _anynew = None

    def __init__(self, nsize: int,
                 maxsize: int,
                 hashfun,
                 heap : list = None,
                 nvisited: int = 0):
        """
        - nsize: size of the ngrams / kmers
        - maxsize: maximum size for the sketch
        - hashfun: function used for hashing - `hashfun(byteslike) -> hash value`
        - heap: heapified list (if unsure about what it is, don't change the default)
        - nvisited: number of kmers visited so far
        """
        self._nsize = nsize
        self._maxsize = maxsize
        self._hashfun = hashfun
        if heap is None:
            self._heap = list()
        else:
            self._heap = heap
        self._heapset = set(self._heap)
        self._nvisited = nvisited

    @property
    def maxsize(self):
        """ Maximum size for the sketch. """
        return self._maxsize

    @property
    def nsize(self):
        """ Size of the ngrams / kmers. """
        return self._nsize

    @property
    def nvisited(self):
        """ Number of ngrams / kmers visited (considered for inclusion) so far. """
        return self._nvisited

    def _add_elt_unsafe(self, elt):
        """
        Add an element.

        - elt: an object as returned by the method `make_elt()`

        Note: This method does not check whether the element is satisfying
        the MaxHash property,
        """
        self._heapset.add(self._extracthash(elt))
        heappush(self._heap, elt)

    def _replace(self, h, elt):
        """
        insert/replace an element

        - h: a hash value
        - elt: an object as returned by the method `make_elt()`
        """
        heapset = self._heapset
        heapset.add(h)
        out = heapreplace(self._heap, elt)
        heapset.remove(self._extracthash(out))
        return out

    @staticmethod
    def _make_elt(h, ngram):
        """
        Make an element to store into the sketch

        - h: a hash value
        - ngram: the object (ngram/kmer) at the source of the hash value
        """
        return (h, ngram)

    @staticmethod
    def _extracthash(heaptop):
        return heaptop[0]

    def __len__(self):
        """
        Return the number of elements in the sketch. See also the property 'nvisited'.
        """
        return len(self._heap)

    def __contains__(self, elt):
        """
        Return whether a given element is in the sketch

        - elt: an object as returned by the method `make_elt()`
        """
        return elt in self._heapset

    def add_hashvalues(self, values):
        """
        Add hash values while conserving the MaxHash characteristic of the set.

        Note: The attribute `nvisited` is not incremented.

        - values: an iterable of hash values
        """
        make_elt = self._make_elt
        anynew = self._anynew
        extracthash = self._extracthash
        heap = self._heap
        heapset = self._heapset
        maxsize = self._maxsize
        lheap = len(heap)
        if lheap > 0:
            heaptop = extracthash(heap[0])
        else:
            heaptop = None
        ngram = None
        for h in values:
            if lheap < maxsize:
                if h not in heapset:
                    elt = make_elt(h, ngram)
                    self._add_elt_unsafe(elt)
                    heaptop = extracthash(heap[0])
                    lheap += 1
                if anynew is not None:
                    anynew(elt)
            if h  >= heaptop:
                if h not in heapset:
                    elt = make_elt(h, ngram)
                    out = self._replace(h, elt)
                    heaptop = extracthash(heap[0])
                if anynew is not None:
                    anynew(elt)

    def _add(self, subs, nsubs, hashbuffer, heaptop,
             extracthash, make_elt, replace, anynew) -> int:
        """
        Process/add elements to the sketch.

        If calling this method directly, updating the attribute
        `nvisited` is under your responsibility.

        - subs: (sub-)sequence
        - nsubs: number of hash values in the hashbuffer
        - hashbuffer: buffer with hash values
        - heaptop: top of the heap
        - extracthash:
        - make_elt:
        - replace:
        - anynew:
        """

        nsize = self._nsize
        heap = self._heap
        lheap = len(heap)
        heapset = self._heapset
        maxsize = self._maxsize

        for j in range(nsubs):
            h = hashbuffer[j]
            if lheap < maxsize:
                if h not in heapset:
                    ngram = subs[j:(j+nsize)]
                    elt = make_elt(h, ngram)
                    self._add_elt_unsafe(elt)
                    heaptop = extracthash(heap[0])
                    lheap += 1
                if anynew is not None:
                    anynew(h)
            elif h  >= heaptop:
                if h not in heapset:
                    ngram = subs[j:(j+nsize)]
                    elt = make_elt(h, ngram)
                    out = replace(h, elt)
                    heaptop = extracthash(heap[0])
                if anynew is not None:
                    anynew(h)
        return heaptop

    def add(self, seq, hashbuffer=array.array('Q', [0,]*250)):
        """ Add all sub-sequences of length `self.nsize` found in the sequence "seq".

        - seq: a bytes-like sequence than can be sliced, and the slices be consummed
               by the function in the property `hashfun` (given to the constructor)
        - hashbuffer: a buffer array to store hash values during batch C calls

        """
        hashfun = self._hashfun
        heap = self._heap
        maxsize = self._maxsize
        nsize = self._nsize
        lseq = len(seq)

        w = len(hashbuffer)
        assert nsize <= w

        anynew = self._anynew
        make_elt = self._make_elt
        extracthash = self._extracthash

        lheap = len(heap)
        if lheap > 0:
            heaptop = extracthash(heap[0])
        else:
            heaptop = 0

        for slice_beg, slice_end in chunkpos_iter(nsize, lseq, w):
            subs = seq[slice_beg:slice_end] # safe: no out-of-bound in Python
            nsubs = hashfun(subs, nsize, hashbuffer)
            heaptop = self._add(subs, nsubs, hashbuffer, heaptop,
                                extracthash, make_elt, self._replace, anynew)
            self._nvisited += nsubs

    def update(self, obj):
        """
        Update the sketch with elements from `obj`.

        - obj: an iterable of elements (each element as returned by `_make_elt()`
        """

        extracthash = self._extracthash
        anynew = self._anynew
        heap = self._heap
        lheap = len(heap)
        heapset = self._heapset
        maxsize = self._maxsize
        if lheap > 0:
            heaptop = extracthash(heap[0])
        else:
            heaptop = 0

        for elt in obj:
            h = extracthash(elt)
            if lheap < maxsize:
                if h not in heapset:
                    self._add_elt_unsafe(elt)
                    heaptop = extracthash(heap[0])
                    lheap += 1
                if anynew is not None:
                    anynew(elt)
            if h  >= heaptop:
                if h not in heapset:
                    out = self._replace(h, elt)
                    heaptop = extracthash(heap[0])
                if anynew is not None:
                    anynew(elt)

        self._nvisited += obj.nvisited

    def __iter__(self):
        """
        Return an iterator over the elements in the sketch.
        """
        return iter(sorted(self._heap))

class MaxHashNgramCountSketch(MaxHashNgramSketch):
    """
    MaxHash Sketch where the number of times an hash value was found is stored
    """

    def __init__(self, nsize: int, maxsize: int, hashfun,
                 heap : list = None,
                 count : Counter = None,
                 nvisited: int = 0):
        """
        - nsize: size of the ngrams / kmers
        - maxsize: maximum size for the sketch
        - hashfun: function used for hashing - `hashfun(byteslike) -> hash value`
        - heap: heapified list (if unsure about what it is, don't change the default)
        - count: a collections.Counter
        - nvisited: number of kmers visited so far
        """
        super().__init__(nsize, maxsize, hashfun,
                         heap = heap, nvisited = nvisited)
        if count is None:
            count = Counter()
            if heap is not None:
                for elt in heap:
                    h = self._extracthash(elt)
                    if h in count:
                        raise ValueError('Elements in the heap must be unique.')
                    else:
                        count[h] = 1
        else:
            if len(self._heapset ^ set(count.keys())) > 0:
                raise ValueError("Mismatching keys with the parameter 'count'.")
        self._count = count

    def _add_elt_unsafe(self, elt):
        super()._add_elt_unsafe(elt)

    def _replace(self, h, elt):
        out = super()._replace(h, elt)
        del(self._count[self._extracthash(out)])
        return out

    def _anynew(self, h):
        self._count[h] += 1

class FrozenHashNgramSketch(object):
    """
    Read-only sketch.
    """

    def __init__(self, sketch : set, nsize : int, maxsize : int = None, nvisited: int = None):
        """
        Create an instance from:
        - setobj: a set
        - nsize: a kmer/ngram size
        - maxsize: a maximum size for the input set (if missing, this is assumed to be len(setobj)
        - nvisited: the number of kmers/ngrams visited to create setobj
        """

        if maxsize is None:
            maxsize = len(sketch)
        elif maxsize < len(sketch):
            raise ValueError("The maximum size cannot be smaller than the number of objects in the set.")
        if nvisited is None:
            nvisited = len(sketch)
        elif nvisited < len(sketch):
            raise ValueError("'nvisited' cannot be smaller than the number of objects in the set.")

        self._sketch = frozenset(sketch)
        self._nsize = nsize
        self._maxsize = maxsize
        self._nvisited = nvisited

    @property
    def maxsize(self):
        """ Maximum size for the sketch. """
        return self._maxsize

    @property
    def nsize(self):
        """ Size of the ngrams / kmers. """
        return self._nsize

    @property
    def nvisited(self):
        """ Number of ngrams / kmers visited (considered for inclusion) so far. """
        return self._nvisited

    def jaccard(self, obj):
        """ Compute the Jaccard index between this sketch and an other sketch"""
        return len(self._sketch.intersection(obj._sketch)) /  len(self._sketch.union(obj._sketch))

    def __len__(self):
        """ Return the number of elements in the set. """
        return len(self._sketch)
<commit_msg>removed unnecessay method in hash sketch-with-counts class<commit_after>from heapq import heappush, heapreplace
from collections import Counter
import array

def chunkpos_iter(nsize: int, lseq: int, w: int) -> (int, int):
    """
    Iterator of chunk indices.

    - nsize: n in ngram
    - lseq: length of sequence to chunk
    - w: width of buffer

    """

    assert nsize <= w
    ew = w-nsize+1

    nchunks = (lseq // ew)

    leftover = lseq - nchunks*ew

    if leftover >= nsize:
        nchunks += 1

    for w_i in range(nchunks):
        slice_beg = (w_i*ew)
        slice_end = slice_beg + w
        yield (slice_beg, slice_end)

class MaxHashNgramSketch(object):

    """
    MaxHash Sketch.
    """

    _anynew = None

    def __init__(self, nsize: int,
                 maxsize: int,
                 hashfun,
                 heap : list = None,
                 nvisited: int = 0):
        """
        - nsize: size of the ngrams / kmers
        - maxsize: maximum size for the sketch
        - hashfun: function used for hashing - `hashfun(byteslike) -> hash value`
        - heap: heapified list (if unsure about what it is, don't change the default)
        - nvisited: number of kmers visited so far
        """
        self._nsize = nsize
        self._maxsize = maxsize
        self._hashfun = hashfun
        if heap is None:
            self._heap = list()
        else:
            self._heap = heap
        self._heapset = set(self._heap)
        self._nvisited = nvisited

    @property
    def maxsize(self):
        """ Maximum size for the sketch. """
        return self._maxsize

    @property
    def nsize(self):
        """ Size of the ngrams / kmers. """
        return self._nsize

    @property
    def nvisited(self):
        """ Number of ngrams / kmers visited (considered for inclusion) so far. """
        return self._nvisited

    def _add_elt_unsafe(self, elt):
        """
        Add an element.

        - elt: an object as returned by the method `make_elt()`

        Note: This method does not check whether the element is satisfying
        the MaxHash property,
        """
        self._heapset.add(self._extracthash(elt))
        heappush(self._heap, elt)

    def _replace(self, h, elt):
        """
        insert/replace an element

        - h: a hash value
        - elt: an object as returned by the method `make_elt()`
        """
        heapset = self._heapset
        heapset.add(h)
        out = heapreplace(self._heap, elt)
        heapset.remove(self._extracthash(out))
        return out

    @staticmethod
    def _make_elt(h, ngram):
        """
        Make an element to store into the sketch

        - h: a hash value
        - ngram: the object (ngram/kmer) at the source of the hash value
        """
        return (h, ngram)

    @staticmethod
    def _extracthash(heaptop):
        return heaptop[0]

    def __len__(self):
        """
        Return the number of elements in the sketch. See also the property 'nvisited'.
        """
        return len(self._heap)

    def __contains__(self, elt):
        """
        Return whether a given element is in the sketch

        - elt: an object as returned by the method `make_elt()`
        """
        return elt in self._heapset

    def add_hashvalues(self, values):
        """
        Add hash values while conserving the MaxHash characteristic of the set.

        Note: The attribute `nvisited` is not incremented.

        - values: an iterable of hash values
        """
        make_elt = self._make_elt
        anynew = self._anynew
        extracthash = self._extracthash
        heap = self._heap
        heapset = self._heapset
        maxsize = self._maxsize
        lheap = len(heap)
        if lheap > 0:
            heaptop = extracthash(heap[0])
        else:
            heaptop = None
        ngram = None
        for h in values:
            if lheap < maxsize:
                if h not in heapset:
                    elt = make_elt(h, ngram)
                    self._add_elt_unsafe(elt)
                    heaptop = extracthash(heap[0])
                    lheap += 1
                if anynew is not None:
                    anynew(elt)
            if h  >= heaptop:
                if h not in heapset:
                    elt = make_elt(h, ngram)
                    out = self._replace(h, elt)
                    heaptop = extracthash(heap[0])
                if anynew is not None:
                    anynew(elt)

    def _add(self, subs, nsubs, hashbuffer, heaptop,
             extracthash, make_elt, replace, anynew) -> int:
        """
        Process/add elements to the sketch.

        If calling this method directly, updating the attribute
        `nvisited` is under your responsibility.

        - subs: (sub-)sequence
        - nsubs: number of hash values in the hashbuffer
        - hashbuffer: buffer with hash values
        - heaptop: top of the heap
        - extracthash:
        - make_elt:
        - replace:
        - anynew:
        """

        nsize = self._nsize
        heap = self._heap
        lheap = len(heap)
        heapset = self._heapset
        maxsize = self._maxsize

        for j in range(nsubs):
            h = hashbuffer[j]
            if lheap < maxsize:
                if h not in heapset:
                    ngram = subs[j:(j+nsize)]
                    elt = make_elt(h, ngram)
                    self._add_elt_unsafe(elt)
                    heaptop = extracthash(heap[0])
                    lheap += 1
                if anynew is not None:
                    anynew(h)
            elif h  >= heaptop:
                if h not in heapset:
                    ngram = subs[j:(j+nsize)]
                    elt = make_elt(h, ngram)
                    out = replace(h, elt)
                    heaptop = extracthash(heap[0])
                if anynew is not None:
                    anynew(h)
        return heaptop

    def add(self, seq, hashbuffer=array.array('Q', [0,]*250)):
        """ Add all sub-sequences of length `self.nsize` found in the sequence "seq".

        - seq: a bytes-like sequence than can be sliced, and the slices be consummed
               by the function in the property `hashfun` (given to the constructor)
        - hashbuffer: a buffer array to store hash values during batch C calls

        """
        hashfun = self._hashfun
        heap = self._heap
        maxsize = self._maxsize
        nsize = self._nsize
        lseq = len(seq)

        w = len(hashbuffer)
        assert nsize <= w

        anynew = self._anynew
        make_elt = self._make_elt
        extracthash = self._extracthash

        lheap = len(heap)
        if lheap > 0:
            heaptop = extracthash(heap[0])
        else:
            heaptop = 0

        for slice_beg, slice_end in chunkpos_iter(nsize, lseq, w):
            subs = seq[slice_beg:slice_end] # safe: no out-of-bound in Python
            nsubs = hashfun(subs, nsize, hashbuffer)
            heaptop = self._add(subs, nsubs, hashbuffer, heaptop,
                                extracthash, make_elt, self._replace, anynew)
            self._nvisited += nsubs

    def update(self, obj):
        """
        Update the sketch with elements from `obj`.

        - obj: an iterable of elements (each element as returned by `_make_elt()`
        """

        extracthash = self._extracthash
        anynew = self._anynew
        heap = self._heap
        lheap = len(heap)
        heapset = self._heapset
        maxsize = self._maxsize
        if lheap > 0:
            heaptop = extracthash(heap[0])
        else:
            heaptop = 0

        for elt in obj:
            h = extracthash(elt)
            if lheap < maxsize:
                if h not in heapset:
                    self._add_elt_unsafe(elt)
                    heaptop = extracthash(heap[0])
                    lheap += 1
                if anynew is not None:
                    anynew(elt)
            if h  >= heaptop:
                if h not in heapset:
                    out = self._replace(h, elt)
                    heaptop = extracthash(heap[0])
                if anynew is not None:
                    anynew(elt)

        self._nvisited += obj.nvisited

    def __iter__(self):
        """
        Return an iterator over the elements in the sketch.
        """
        return iter(sorted(self._heap))

class MaxHashNgramCountSketch(MaxHashNgramSketch):
    """
    MaxHash Sketch where the number of times an hash value was found is stored
    """

    def __init__(self, nsize: int, maxsize: int, hashfun,
                 heap : list = None,
                 count : Counter = None,
                 nvisited: int = 0):
        """
        - nsize: size of the ngrams / kmers
        - maxsize: maximum size for the sketch
        - hashfun: function used for hashing - `hashfun(byteslike) -> hash value`
        - heap: heapified list (if unsure about what it is, don't change the default)
        - count: a collections.Counter
        - nvisited: number of kmers visited so far
        """
        super().__init__(nsize, maxsize, hashfun,
                         heap = heap, nvisited = nvisited)
        if count is None:
            count = Counter()
            if heap is not None:
                for elt in heap:
                    h = self._extracthash(elt)
                    if h in count:
                        raise ValueError('Elements in the heap must be unique.')
                    else:
                        count[h] = 1
        else:
            if len(self._heapset ^ set(count.keys())) > 0:
                raise ValueError("Mismatching keys with the parameter 'count'.")
        self._count = count

    def _replace(self, h, elt):
        out = super()._replace(h, elt)
        del(self._count[self._extracthash(out)])
        return out

    def _anynew(self, h):
        self._count[h] += 1

class FrozenHashNgramSketch(object):
    """
    Read-only sketch.
    """

    def __init__(self, sketch : set, nsize : int, maxsize : int = None, nvisited: int = None):
        """
        Create an instance from:
        - setobj: a set
        - nsize: a kmer/ngram size
        - maxsize: a maximum size for the input set (if missing, this is assumed to be len(setobj)
        - nvisited: the number of kmers/ngrams visited to create setobj
        """

        if maxsize is None:
            maxsize = len(sketch)
        elif maxsize < len(sketch):
            raise ValueError("The maximum size cannot be smaller than the number of objects in the set.")
        if nvisited is None:
            nvisited = len(sketch)
        elif nvisited < len(sketch):
            raise ValueError("'nvisited' cannot be smaller than the number of objects in the set.")

        self._sketch = frozenset(sketch)
        self._nsize = nsize
        self._maxsize = maxsize
        self._nvisited = nvisited

    @property
    def maxsize(self):
        """ Maximum size for the sketch. """
        return self._maxsize

    @property
    def nsize(self):
        """ Size of the ngrams / kmers. """
        return self._nsize

    @property
    def nvisited(self):
        """ Number of ngrams / kmers visited (considered for inclusion) so far. """
        return self._nvisited

    def jaccard(self, obj):
        """ Compute the Jaccard index between this sketch and an other sketch"""
        return len(self._sketch.intersection(obj._sketch)) /  len(self._sketch.union(obj._sketch))

    def __len__(self):
        """ Return the number of elements in the set. """
        return len(self._sketch)
<|endoftext|>


<commit_before>import sys
import rtmidi
import threading

def print_message(midi, port):
    if midi.isNoteOn():
        print '%s: ON: ' % port, \
            midi.getMidiNoteName(midi.getNoteNumber()), \
            midi.getVelocity()
    elif midi.isNoteOff():
        print '%s: OFF:' % port, \
            midi.getMidiNoteName(midi.getNoteNumber())
    elif midi.isController():
        print '%s: CONTROLLER' % port, \
            midi.getControllerNumber(), midi.getControllerValue()

def add_event(midi, port):
    pass

class Collector(threading.Thread):
    def __init__(self, device, port):
        threading.Thread.__init__(self)
        self.setDaemon(True)
        self.port = port
        self.portName = device.getPortName(port)
        self.device = device
        self.quit = False

    def run(self):
        self.device.openPort(self.port)
        self.device.ignoreTypes(True, False, True)
        while True:
            if self.quit:
                return
            msg = self.device.getMessage()
            if msg:
                print_message(msg, self.portName)

dev = rtmidi.RtMidiIn()
collectors = []
for i in range(dev.getPortCount()):
    device = rtmidi.RtMidiIn()
    print 'OPENING', dev.getPortName(i)
    collector = Collector(device, i)
    collector.start()
    collectors.append(collector)

print 'HIT ENTER TO EXIT'
sys.stdin.read(1)
for c in collectors:
    c.quit = True
<commit_msg>Reduced performance impact by waiting<commit_after>import sys
import time
import rtmidi
import threading

def print_message(midi, port):
    if midi.isNoteOn():
        print '%s: ON: ' % port, \
            midi.getMidiNoteName(midi.getNoteNumber()), \
            midi.getVelocity()
    elif midi.isNoteOff():
        print '%s: OFF:' % port, \
            midi.getMidiNoteName(midi.getNoteNumber())
    elif midi.isController():
        print '%s: CONTROLLER' % port, \
            midi.getControllerNumber(), midi.getControllerValue()

def add_event(midi, port):
    pass

class Collector(threading.Thread):
    def __init__(self, device, port):
        threading.Thread.__init__(self)
        self.setDaemon(True)
        self.port = port
        self.portName = device.getPortName(port)
        self.device = device
        self.quit = False

    def run(self):
        self.device.openPort(self.port)
        self.device.ignoreTypes(True, False, True)
        while True:
            if self.quit:
                return
            msg = self.device.getMessage()
            if msg:
                print_message(msg, self.portName)
            else:
                time.sleep(0.1)

dev = rtmidi.RtMidiIn()
collectors = []
for i in range(dev.getPortCount()):
    device = rtmidi.RtMidiIn()
    print 'OPENING', dev.getPortName(i)
    collector = Collector(device, i)
    collector.start()
    collectors.append(collector)

print 'HIT ENTER TO EXIT'
sys.stdin.read(1)
for c in collectors:
    c.quit = True
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# shredcompare -- find similarities in file trees.
#
# 1.0 by Eric S. Raymond, 24 Aug 2003.
#
# Implements the shred algorithm described at:
#
#	http://theinquirer.net/?article=10061
#
#   "Here is the procedure for finding the matching code.... "1. Each
#   file withing each source tree is "shredded" into 5 line pieces (1-5,
#   2-6, 3-7, etc.). MD5 sum is computed for each block of lines. The
#   output is 3 columns: MD5sum, source file, 1st line in the block.
#
#   "At this stage, 4.4BSD had [a] ~40Mb file, linux ~160Mb. Potentially,
#   one could shred into smaller or larger pieces, however, with pieces
#   too small there'll be a lot of noise, with pieces too large some
#   matches won't be seen. 5 liners seem to be a good compromise.
#
#   "2. Within each source tree the "shredded" file is sorted by MD5sum,
#   and duplicate entries within the same tree are removed completely
#   (these are either trivial 5-line sequences or licensing
#   disclaimers). Unix sort here takes a couple of minutes on a 600Mhz
#   P3.
#
#   "3. A column indicating the origin of the file is inserted into the
#   file (0 - BSD, 1 - linux). Both Linux and BSD "shredded" files are
#   merged such that MD5sums stay sorted.
#
#   "4. At this point a given MD5sum will occur either once or twice,
#   i.e., in both source trees. Here remove all thesingle lines, and have
#   the 5 liners left that are matching.
#
#   "5. Count for each file in Linux tree the number of matches with the
#   BSD tree using the file generated at step 4. Sort this list, and the
#   largest counts will occur for the files with the largest number of
#   matching lines. The range can be extracted from the file from step 4,
#   since at step 1 we kept the address of the 1st line in the
#   block. That is how the info above was generated.
#
# In this implementation, lines are preprocessed to remove whitespace
# differences.
#
# The -h option supports building a cache containing the shred list
# for a tree you specify.  This is strictly a speed hack.  If your
# tree is named `foo', the file will be named `foo.hash'.  On your
# next comparison run it will be picked up automatically. It is
# your responsibility to make sure each hash file is newer than
# the correspobding tree.

import sys, os, os.path, re, md5, getopt, time, cPickle

shredsize = 5
verbose = False
mark_time = None
ws = re.compile(r"\s+")

class Range:
    "Represent a range of lines."
    def __init__(self, file, start, end):
        self.file = file
        self.start = start
        self.end = end
    def fetch(self):
        rfp = open(self.file)
        for i in range(self.start):
            rfp.readline()
        text = ""
        for i in range(self.start, self.end):
            text += rfp.readline()
        rfp.close()
        return text
    def __cmp__(self, other):
        return cmp((self.file, self.start, self.end),
                   (other.file, other.start, other.end))
    def __str__(self):
        # 0-origin line numbers internally, 1-origin externally.
        return "%s:%d-%d" % (self.file, self.start+1, self.end)
    __repr__ = __str__

def smash_whitespace(line):
    "Replace each string of whitespace in a line with a single space."
    return ws.sub(' ', line)

def shredfile(file, tree):
    "Build a dictionary of shred tuples corresponding to a specified file."
    if verbose:
        print file
    shreds = {}
    fp = open(file)
    lines = map(smash_whitespace, fp.readlines())
    fp.close()
    for i in range(len(lines)-shredsize):
        m = md5.new()
        for j in range(shredsize):
            m.update(lines[i+j])
        # Merging shreds into a dict automatically suppresses duplicates
        shreds[m.digest()] = Range(file, i, i+shredsize)
    return shreds

def eligible(file):
    "Is a file eligible for comparison?"
    return filter(lambda x: file.endswith(x), ('.c', 'h', '.y', '.l', '.txt'))

def shredtree(tree):
    "Shred an entire file tree; return a dictionary of checksums-to-locations."
    shreds = {}
    if os.path.exists(tree + ".hash"):
        # There's a precomputed shred list.
        rfp = open(tree + ".hash")
        shreds = cPickle.load(rfp)
        rfp.close()
    else:
        # No precomputed hashes.  Prepare a news shred list.
        for root, dirs, files in os.walk(tree):
            for file in files:
                path = os.path.join(root, file)
                if eligible(path):
                    shreds.update(shredfile(path, tree))
    return shreds

def shredcompare(tree1, tree2):
    "Compare two trees.  Returns a list of matching Range pairs"
    shreds1 = shredtree(tree1)
    shreds2 = shredtree(tree2)
    # Nuke everything but checksum matches between the trees
    for key in shreds1.keys():
        if not key in shreds2:
            del shreds1[key]
    for key in shreds2.keys():
        if not key in shreds1:
            del shreds2[key]
    # Merge the match lists
    matches = []
    for key in shreds1.keys():
        matches.append((shreds1[key], shreds2[key]))
    # Merge adjacent ranges if possible
    matches.sort()	# Sort in natural order
    merge_hit = True
    while merge_hit:
        merge_hit = False
        j = len(matches) - 1
        while j > 0:
            this = matches[j]
            last = matches[j-1]
            if last[0].file == this[0].file \
		   and last[1].file==this[1].file \
                   and last[0].start==this[0].start-1:
                merge_hit = True
                last[0].end = this[0].end
                last[1].end = this[1].end
                matches = matches[:j] + matches[j+1:]
            j -= 1
    return matches

def report_time(legend=None):
    "Report time since start_time was set."
    global mark_time
    endtime = time.time()
    if mark_time and verbose:
        elapsed = endtime - mark_time
        hours = elapsed/3600; elapsed %= 3600
        minutes = elapsed/60; elapsed %= 60
        seconds = elapsed
        print "%% %s: %dh, %dm, %ds" % (legend, hours, minutes, seconds)
    mark_time = endtime

if __name__ == '__main__':
    try:
        (optlist, args) = getopt.getopt(sys.argv[1:], 'h:s:v')
    except getopt.GetoptError:
        sys.stderr.write("usage: shredcompare [-s shredsize] [-v] [-h] tree1 tree2\n")
        sys.exit(2)
    makehash = None
    for (opt, val) in optlist:
        if opt == '-s':
            shredsize = int(val)
        elif opt == '-v':
            verbose = True
        elif opt == '-h':
            makehash = val
    report_time()
    if makehash:
        shreds = shredtree(makehash)
        wfp = open(makehash + ".hash", "w")
        report_time("Shred list complete.")
        cPickle.dump(shreds, wfp)
        wfp.close()
        report_time("Pickling cmplete")
    else:
        matches = shredcompare(args[0], args[1])
        report_time("Match list complete")
        for (source, target) in matches:
            print source, "->", target
            if verbose:
                sys.stdout.write(source.fetch())

<commit_msg>Better comments.<commit_after>#!/usr/bin/env python
#
# shredcompare -- find similarities in file trees.
#
# 0.1 by Eric S. Raymond, 24 Aug 2003.
#
# Implements the shred algorithm described at:
#
#	http://theinquirer.net/?article=10061
#
# The -h option supports building a cache containing the shred list
# for a tree you specify.  This is strictly a speed hack.  If your
# tree is named `foo', the file will be named `foo.hash'.  On your
# next comparison run it will be picked up automatically. It is
# your responsibility to make sure each hash file is newer than
# the corresponding tree.
#
# This is alpha software. There is a bug in the implementation,
# somewhere in the logic for merging matching ranges before report
# generation, that sometimes garbages the last number in the
# target-tree range.  I haven't fixed it yet because I haven't found a
# way to produce a test load that reproduces it that is smaller that
# the entire Linux 2.6.o and Sysem V kernels.

import sys, os, os.path, re, md5, getopt, time, cPickle

shredsize = 5
verbose = False
mark_time = None
ws = re.compile(r"\s+")

class Range:
    "Represent a range of lines."
    def __init__(self, file, start, end):
        self.file = file
        self.start = start
        self.end = end
    def fetch(self):
        rfp = open(self.file)
        for i in range(self.start):
            rfp.readline()
        text = ""
        for i in range(self.start, self.end):
            text += rfp.readline()
        rfp.close()
        return text
    def __cmp__(self, other):
        return cmp((self.file, self.start, self.end),
                   (other.file, other.start, other.end))
    def __str__(self):
        # 0-origin line numbers internally, 1-origin externally.
        return "%s:%d-%d" % (self.file, self.start+1, self.end)
    __repr__ = __str__

def smash_whitespace(line):
    "Replace each string of whitespace in a line with a single space."
    return ws.sub(' ', line)

def shredfile(file, tree):
    "Build a dictionary of shred tuples corresponding to a specified file."
    if verbose:
        print file
    shreds = {}
    fp = open(file)
    lines = map(smash_whitespace, fp.readlines())
    fp.close()
    for i in range(len(lines)-shredsize):
        m = md5.new()
        for j in range(shredsize):
            m.update(lines[i+j])
        # Merging shreds into a dict automatically suppresses duplicates
        shreds[m.digest()] = Range(file, i, i+shredsize)
    return shreds

def eligible(file):
    "Is a file eligible for comparison?"
    return filter(lambda x: file.endswith(x), ('.c', 'h', '.y', '.l', '.txt'))

def shredtree(tree):
    "Shred an entire file tree; return a dictionary of checksums-to-locations."
    shreds = {}
    if os.path.exists(tree + ".hash"):
        # There's a precomputed shred list.
        rfp = open(tree + ".hash")
        shreds = cPickle.load(rfp)
        rfp.close()
    else:
        # No precomputed hashes.  Prepare a news shred list.
        for root, dirs, files in os.walk(tree):
            for file in files:
                path = os.path.join(root, file)
                if eligible(path):
                    shreds.update(shredfile(path, tree))
    return shreds

def shredcompare(tree1, tree2):
    "Compare two trees.  Returns a list of matching Range pairs"
    shreds1 = shredtree(tree1)
    shreds2 = shredtree(tree2)
    # Nuke everything but checksum matches between the trees
    for key in shreds1.keys():
        if not key in shreds2:
            del shreds1[key]
    for key in shreds2.keys():
        if not key in shreds1:
            del shreds2[key]
    # Merge the match lists
    matches = []
    for key in shreds1.keys():
        matches.append((shreds1[key], shreds2[key]))
    # Merge adjacent ranges if possible
    matches.sort()	# Sort in natural order
    merge_hit = True
    while merge_hit:
        merge_hit = False
        j = len(matches) - 1
        while j > 0:
            this = matches[j]
            last = matches[j-1]
            if last[0].file == this[0].file \
		   and last[1].file==this[1].file \
                   and last[0].start==this[0].start-1:
                merge_hit = True
                last[0].end = this[0].end
                last[1].end = this[1].end
                matches = matches[:j] + matches[j+1:]
            j -= 1
    return matches

def report_time(legend=None):
    "Report time since start_time was set."
    global mark_time
    endtime = time.time()
    if mark_time and verbose:
        elapsed = endtime - mark_time
        hours = elapsed/3600; elapsed %= 3600
        minutes = elapsed/60; elapsed %= 60
        seconds = elapsed
        print "%% %s: %dh, %dm, %ds" % (legend, hours, minutes, seconds)
    mark_time = endtime

if __name__ == '__main__':
    try:
        (optlist, args) = getopt.getopt(sys.argv[1:], 'h:s:v')
    except getopt.GetoptError:
        sys.stderr.write("usage: shredcompare [-s shredsize] [-v] [-h] tree1 tree2\n")
        sys.exit(2)
    makehash = None
    for (opt, val) in optlist:
        if opt == '-s':
            shredsize = int(val)
        elif opt == '-v':
            verbose = True
        elif opt == '-h':
            makehash = val
    report_time()
    if makehash:
        shreds = shredtree(makehash)
        wfp = open(makehash + ".hash", "w")
        report_time("Shred list complete.")
        cPickle.dump(shreds, wfp)
        wfp.close()
        report_time("Pickling complete")
    else:
        matches = shredcompare(args[0], args[1])
        report_time("Match list complete")
        for (source, target) in matches:
            print source, "->", target
            if verbose:
                sys.stdout.write(source.fetch())

<|endoftext|>


<commit_before># !/usr/bin/env python
"""Set Python environment."""
# -*- coding: utf-8 -*-

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

from splinter import Browser
from fuzzywuzzy import fuzz

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:
    """Object for testing lists of XSS payloads."""

    def __init__(self):
        """Initiate object with default encoding and other required data."""
        # Expect some weird characters from fuzz lists, make encoding UTF-8
        reload(sys)
        sys.setdefaultencoding('utf8')

        # All potential XSS findings
        self.xss_links = []

        # Fuzzy string matching partials list
        self.xss_partials = []

        # Keep index of screens for log files
        self.screen_index = ""

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

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

    def main(self):
        """Call functions to inject XSS and test for vulnerabilities."""
        # 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 + \
                "XSS vulnerabilities were detected!" + \
                Color.END
            self.log_file(self.xss_links)
        # There were partials detected by fuzzy detection
        elif self.xss_partials:
            print Color.YELLOW + \
                "Partial 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):
        """Ensure that file path exists before writing."""
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise

    def inject_payload(self, payload, link, request_delay,
                       user_screenshot_name):
        """Inject XSS payload string from user supplied payload list."""
        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 so, take screenshot depending on user flag
        self.detect_xss(payload, browser, user_screenshot_name, injected_link)

    def detect_xss(self, payload, browser_object, user_screenshot_name,
                   injected_link):
        """Check the HTML source to determine if XSS payload was reflected."""
        # If fuzzy detection chosen, evaluate partial reflection of XSS
        # by tokenizing the HTML source and detecting parts of the payload
        # and source common to both.
        #
        # Other methods of scoring include fuzz.ratio(), fuzz.partial_ratio()
        # and fuzz.token_sort_ratio()
        partial_score = fuzz.token_set_ratio(
            payload.lower(), browser_object.html.lower())
        # Set the level of detection asked for by the user, e.g. Only detect
        # matches with score higher than 50% fuzzy detection
        fuzzy_level = self.user_args.FUZZY_DETECTION

        if payload.lower() in browser_object.html.lower():
            print Color.GREEN + "\n[+] XSS vulnerability found:" + \
                Color.END

            # If user set the --screen flag to target, capture screenshot of
            # payload
            if user_screenshot_name is not None:
                self.take_screenshot(user_screenshot_name,
                                     browser_object, self.screen_index)

            # Add link to list of all positive XSS hits
            self.xss_links.append(injected_link)
            print Color.BLUE + injected_link + Color.END
        # If user enabled fuzzy detection and partial score was larger than
        # fuzz level, add it to partials list and print results
        elif fuzzy_level and (partial_score >= fuzzy_level):
            print Color.YELLOW + \
                "\n[-] Partial XSS vulnerability found:" + Color.END
            print Color.BLUE + injected_link + Color.END
            self.xss_partials.append(injected_link)
            print "Detection score: %s" % partial_score
        else:
            print Color.RED + "\n[+] No XSS detected at: \n" + \
                Color.BLUE + injected_link + Color.END
            if (fuzzy_level):
                print "Detection score: %s" % partial_score

    def take_screenshot(self, user_screenshot_name, browser_object,
                        screen_index):
        """Take a screenshot of the page in the browser object."""
        # Check if screenshots directory exists, if not then create it
        self.make_sure_path_exists("screenshots")
        screenshot_file_name = "screenshots/" + \
            user_screenshot_name + "_" + \
            datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + \
            "_" + 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,
                 user_screenshot_name):
        """Load string from payload list and call function to inject it."""
        # 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:
            self.inject_payload(item, link, request_delay,
                                user_screenshot_name)

    def log_file(self, link_list):
        """Log successful XSS payload reflections to file."""
        # 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")
            # Set file name
            file_name = "logs/" + target_name + "_" + \
                datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            # Save log file to directory if XSS links is populated
            if self.xss_links:
                with open(file_name + ".txt", 'w') as link_file:
                    for link in link_list:
                        link_file.write(link)
                        link_file.write("\n")
                    # Add metadata about what payload file was used
                    link_file.write(
                        "\n*** Created from the payload file >>> " +
                        self.user_args.PAYLOADS_LIST)
                    link_file.close()
            # If the user enabled fuzzy detection, log partial results
            if (self.user_args.FUZZY_DETECTION):
                with open(file_name + "_partials.txt", 'w') as partial_file:
                    for link in self.xss_partials:
                        partial_file.write(link)
                        partial_file.write("\n")
                    # Add metadata about what payload file was used
                    partial_file.write(
                        "\n*** Created from the payload file >>> " +
                        self.user_args.PAYLOADS_LIST)
                    partial_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 inject XSS payloads into.', required=True)
        parser.add_argument(
            '-p', action='store', dest='PAYLOADS_LIST',
            help='The payload list to use for injection.', 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='Enable screenshots of XSS hits.')
        parser.add_argument(
            '-f', "--fuzzy", action='store', dest='FUZZY_DETECTION', type=int,
            const=50, nargs="?",
            help='Fuzzy detection rate of XSS [0 to 100% match] (default=50).')

        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>Fix arparse error, remove percent sign in help<commit_after># !/usr/bin/env python
"""Set Python environment."""
# -*- coding: utf-8 -*-

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

from splinter import Browser
from fuzzywuzzy import fuzz

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:
    """Object for testing lists of XSS payloads."""

    def __init__(self):
        """Initiate object with default encoding and other required data."""
        # Expect some weird characters from fuzz lists, make encoding UTF-8
        reload(sys)
        sys.setdefaultencoding('utf8')

        # All potential XSS findings
        self.xss_links = []

        # Fuzzy string matching partials list
        self.xss_partials = []

        # Keep index of screens for log files
        self.screen_index = ""

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

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

    def main(self):
        """Call functions to inject XSS and test for vulnerabilities."""
        # 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 + \
                "XSS vulnerabilities were detected!" + \
                Color.END
            self.log_file(self.xss_links)
        # There were partials detected by fuzzy detection
        elif self.xss_partials:
            print Color.YELLOW + \
                "Partial 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):
        """Ensure that file path exists before writing."""
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise

    def inject_payload(self, payload, link, request_delay,
                       user_screenshot_name):
        """Inject XSS payload string from user supplied payload list."""
        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 so, take screenshot depending on user flag
        self.detect_xss(payload, browser, user_screenshot_name, injected_link)

    def detect_xss(self, payload, browser_object, user_screenshot_name,
                   injected_link):
        """Check the HTML source to determine if XSS payload was reflected."""
        # If fuzzy detection chosen, evaluate partial reflection of XSS
        # by tokenizing the HTML source and detecting parts of the payload
        # and source common to both.
        #
        # Other methods of scoring include fuzz.ratio(), fuzz.partial_ratio()
        # and fuzz.token_sort_ratio()
        partial_score = fuzz.token_set_ratio(
            payload.lower(), browser_object.html.lower())
        # Set the level of detection asked for by the user, e.g. Only detect
        # matches with score higher than 50% fuzzy detection
        fuzzy_level = self.user_args.FUZZY_DETECTION

        if payload.lower() in browser_object.html.lower():
            print Color.GREEN + "\n[+] XSS vulnerability found:" + \
                Color.END

            # If user set the --screen flag to target, capture screenshot of
            # payload
            if user_screenshot_name is not None:
                self.take_screenshot(user_screenshot_name,
                                     browser_object, self.screen_index)

            # Add link to list of all positive XSS hits
            self.xss_links.append(injected_link)
            print Color.BLUE + injected_link + Color.END
        # If user enabled fuzzy detection and partial score was larger than
        # fuzz level, add it to partials list and print results
        elif fuzzy_level and (partial_score >= fuzzy_level):
            print Color.YELLOW + \
                "\n[-] Partial XSS vulnerability found:" + Color.END
            print Color.BLUE + injected_link + Color.END
            self.xss_partials.append(injected_link)
            print "Detection score: %s" % partial_score
        else:
            print Color.RED + "\n[+] No XSS detected at: \n" + \
                Color.BLUE + injected_link + Color.END
            if (fuzzy_level):
                print "Detection score: %s" % partial_score

    def take_screenshot(self, user_screenshot_name, browser_object,
                        screen_index):
        """Take a screenshot of the page in the browser object."""
        # Check if screenshots directory exists, if not then create it
        self.make_sure_path_exists("screenshots")
        screenshot_file_name = "screenshots/" + \
            user_screenshot_name + "_" + \
            datetime.datetime.now().strftime("%Y%m%d-%H%M%S") + \
            "_" + 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,
                 user_screenshot_name):
        """Load string from payload list and call function to inject it."""
        # 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:
            self.inject_payload(item, link, request_delay,
                                user_screenshot_name)

    def log_file(self, link_list):
        """Log successful XSS payload reflections to file."""
        # 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")
            # Set file name
            file_name = "logs/" + target_name + "_" + \
                datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            # Save log file to directory if XSS links is populated
            if self.xss_links:
                with open(file_name + ".txt", 'w') as link_file:
                    for link in link_list:
                        link_file.write(link)
                        link_file.write("\n")
                    # Add metadata about what payload file was used
                    link_file.write(
                        "\n*** Created from the payload file >>> " +
                        self.user_args.PAYLOADS_LIST)
                    link_file.close()
            # If the user enabled fuzzy detection, log partial results
            if (self.user_args.FUZZY_DETECTION):
                with open(file_name + "_partials.txt", 'w') as partial_file:
                    for link in self.xss_partials:
                        partial_file.write(link)
                        partial_file.write("\n")
                    # Add metadata about what payload file was used
                    partial_file.write(
                        "\n*** Created from the payload file >>> " +
                        self.user_args.PAYLOADS_LIST)
                    partial_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 inject XSS payloads into.', required=True)
        parser.add_argument(
            '-p', action='store', dest='PAYLOADS_LIST',
            help='The payload list to use for injection.', 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='Enable screenshots of XSS hits.')
        parser.add_argument(
            '-f', "--fuzzy", action='store', dest='FUZZY_DETECTION', type=int,
            const=50, nargs="?",
            help='Fuzzy detection rate of XSS [0 to 100 match] (default=50).')

        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|>


<commit_before>"""mist.io views"""
import json
import logging

from pyramid.response import Response
from libcloud.compute.base import Node, NodeSize, NodeImage, NodeLocation
from libcloud.compute.providers import get_driver
#from libcloud.compute.base import NodeAuthSSHKey
#from libcloud.compute.deployment import MultiStepDeployment, ScriptDeployment
from libcloud.compute.deployment import SSHKeyDeployment
from libcloud.compute.types import Provider
from mist.io.config import BACKENDS, BASE_EC2_AMIS
from pyramid.view import view_config
#from fabric.state import env
#from fabric.api import run
from mist.io.machinecaps import get_machine_actions

LOG = logging.getLogger('mist.io')

def connect(request):
    """Establish backend connection using the credentials specified"""
    try:
        backend_list = request.environ['beaker.session']['backends']
    except:
        backend_list = BACKENDS

    backend_index = int(request.matchdict['backend'])
    backend = backend_list[backend_index]

    driver = get_driver(int(backend['provider']))
    if 'host' in backend.keys():
        conn = driver(backend['id'],
                      backend['secret'],
                      False,
                      host=backend['host'],
                      ex_force_auth_url=backend.get('auth_url', None),
                      ex_force_auth_version=backend.get('auth_version', '1.0'),
                      port=80)
    else:
        conn = driver(backend['id'], backend['secret'])
    return conn

@view_config(route_name='home',
             request_method='GET',
             renderer='templates/home.pt')
def home(request):
    """Fill in an object with backend data, taken from config.py"""
    try:
        backend_list = request.environ['beaker.session']['backends']
        session = True
    except:
        backend_list = BACKENDS
        session = False

    backends = []
    for backend in backend_list:
        backends.append({'id'           : backend['id'],
                         'title'        : backend['title'],
                         'provider'     : backend['provider'],
                         'poll_interval': backend['poll_interval'],
                         'status'       : 'off',
                         })

    return {'project': 'mist.io',
            'backends': backends,
            'session': session}

@view_config(route_name='machines', request_method='GET', renderer='json')
def list_machines(request):
    """List machines for a backend"""

    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        machines = conn.list_nodes()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for m in machines:
        # for rackspace get the tags stored in extra.metadata.tags attr, for amazon get extra.tags.tags attr
        tags = m.extra.get('tags', None) or m.extra.get('metadata', None)
        tags = tags and tags.get('tags', None) or []
        machine = {'id'           : m.id,
                  'uuid'          : m.get_uuid(),
                  'name'          : m.name,
                  # both rackspace and amazon have the image in the imageId extra attr,
                  'imageId'       : m.image or m.extra.get('imageId', None),
                  # for rackspace get flavorId extra attr, for amazon the instancetype extra attr
                  'size'          : m.size or m.extra.get('flavorId', None) or m.extra.get('instancetype', None),
                  'state'         : m.state,
                  'private_ips'   : m.private_ips,
                  'public_ips'    : m.public_ips,
                  'tags'          : tags,
                  'extra'         : m.extra,
                  }
        machine.update(get_machine_actions(m, conn))
        ret.append(machine)
    return ret

@view_config(route_name='machines', request_method='POST', renderer='json')
def create_machine(request):
    """Create a new virtual machine on the specified backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        machine_name = request.json_body['name']
        location_id = request.json_body['location']
        image_id = request.json_body['image']
        size_id = request.json_body['size']
    except Exception as e:
        return Response('Invalid payload', 400)

    size = NodeSize(size_id, name='', ram='', disk='', bandwidth='', price='',
                    driver=conn)
    image = NodeImage(image_id, name='', driver=conn)
    location = NodeLocation(location_id, name='', country='', driver=conn)

    if conn.type == Provider.RACKSPACE and len(request.registry.settings['keypairs']):
        # try to deploy node with ssh key installed
        ssh_key = SSHKeyDeployment(request.registry.settings['keypairs'][0][0])
        try:
            conn.deploy_node(name=machine_name,
                             image=image,
                             size=size,
                             location=location,
                             deploy=ssh_key)
            return []
        except:
            LOG.warn('Failed to deploy node with ssh key. Trying to create simple node')

    try:
        conn.create_node(name=machine_name,
                         image=image,
                         size=size,
                         location=location)
        return []
    except Exception as e:
        return Response('Something went wrong with the creation', 500)

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=start',
             renderer='json')
def start_machine(request):
    """Starts a machine on backends that supported start/stop"""
    #TODO: in which providers could this work?
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    #machine.start()

    return []

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=stop',
             renderer='json')
def stop_machine(request):
    """Stops a machine on backends that supported start/stop"""
    #TODO: in which providers could this work?
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    #machine.stop()

    return []

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=reboot',
             renderer='json')
def reboot_machine(request):
    """Reboots a machine on a certain backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    machine.reboot()
    return []

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=destroy',
             renderer='json')
def destroy_machine(request):
    """Destroys a machine on a certain backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    machine.destroy()

    return []

@view_config(route_name='metadata', request_method='POST')
def set_metadata(request):
    """Sets metadata for a machine, given the backend and machine id"""
    #TODO: the following are not working
    """Examples
    Openstack:
        conn.ex_set_metadata(machine,
                            {'name': 'ServerX',
                             'description': 'all the money'})
    EC2:
        conn2.ex_create_tags(machine, {'something': 'something_something'})
    """
    ret = []
    done = False
    backends = [b for b in BACKENDS if b['id'] == request.matchdict['backend']]
    if backends:
        backend = backends[0]
        conn = connect(backend)
        machines = conn.list_nodes()
        for machine in machines:
            if machine.id == request.matchdict['machine']:
                try:
                    metadata = request.json_body
                    #get metadata from request
                except:
                    return Response('Not proper format for metadata', 404)
                try:
                    #e.g. Openstack
                    metadata = conn.ex_set_metadata(machine, metadata)
                    done = True
                except:
                    try:
                        #e.g. EC2
                        metadata = conn.ex_create_tags(machine, metadata)
                        done = True
                    except:
                        return Response('Not implemented for this backend', 404)
                break
    if not done:
        return Response('Invalid backend', 404)

    return Response(json.dumps(ret))

@view_config(route_name='images', request_method='GET', renderer='json')
def list_images(request):
    """List images from each backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        try:
            backend_list = request.environ['beaker.session']['backends']
        except:
            backend_list = BACKENDS
        backend_index = int(request.matchdict['backend'])
        backend = backend_list[backend_index]
        if backend['provider'] == Provider.EC2:
            images = conn.list_images(None, BASE_EC2_AMIS.keys())
        else:
            images = conn.list_images()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for image in images:
        ret.append({'id'    : image.id,
                    'extra' : image.extra,
                    'name'  : image.name,
                    })
    return ret

@view_config(route_name='sizes', request_method='GET', renderer='json')
def list_sizes(request):
    """List sizes (aka flavors) from each backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        sizes = conn.list_sizes()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for size in sizes:
        ret.append({'id'        : size.id,
                    'bandwidth' : size.bandwidth,
                    'disk'      : size.disk,
                    'driver'    : size.driver.name,
                    'name'      : size.name,
                    'price'     : size.price,
                    'ram'       : size.ram,
                    })

    return ret

@view_config(route_name='locations', request_method='GET', renderer='json')
def list_locations(request):
    """List locations from each backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        locations = conn.list_locations()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for location in locations:
        ret.append({'id'        : location.id,
                    'name'      : location.name,
                    'country'   : location.country,
                    })

    return ret

@view_config(route_name='image_details', request_method='GET', renderer='json')
def get_image_details(request):
    """get_image_details"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        image = conn.get_image(request.params['id'])
    except:
        return Response('Backend unavailable', 503)

    if image is None:
        ret = {}
    else:
        ret = {'id'     : image.id,
               'extra' : image.extra,
               'name'  : image.name,
               }
    return ret

@view_config(route_name='backends', request_method='GET', renderer='json')
def get_backends(request):
    """get_backends"""
    try:
        backend_list = request.environ['beaker.session']['backends']
    except:
        backend_list = BACKENDS

    backends = []
    index = 0
    for backend in backend_list:
        backends.append({'index'        : index,
                         'id'           : backend['id'],
                         'title'        : backend['title'],
                         'provider'     : backend['provider'],
                         'poll_interval': backend['poll_interval'],
                         'status'       : 'online',
                         })
        index = index + 1

    return backends

@view_config(route_name='machine_has_key', request_method='GET', renderer='json')
def machine_has_key(request):
    """has an ssh key been set for this machine"""
    machine_ip = request.params.get('ip', None)

    """ Text code
    env.host_name = machine_ip

    #TODO setup ssh here

    if run('uptime').failed:
        ret = {'has_key': False}
    return ret
    """

    return False
<commit_msg>Test if the machine has ssh key<commit_after>"""mist.io views"""
import os
import logging
import json
import tempfile

from pyramid.response import Response
from libcloud.compute.base import Node, NodeSize, NodeImage, NodeLocation
from libcloud.compute.providers import get_driver
#from libcloud.compute.base import NodeAuthSSHKey
#from libcloud.compute.deployment import MultiStepDeployment, ScriptDeployment
from libcloud.compute.deployment import SSHKeyDeployment
from libcloud.compute.types import Provider
from mist.io.config import BACKENDS, BASE_EC2_AMIS
from pyramid.view import view_config
from fabric.api import run, env
from mist.io.machinecaps import get_machine_actions

log = logging.getLogger('mist.io')

def connect(request):
    """Establish backend connection using the credentials specified"""
    try:
        backend_list = request.environ['beaker.session']['backends']
    except:
        backend_list = BACKENDS

    backend_index = int(request.matchdict['backend'])
    backend = backend_list[backend_index]

    driver = get_driver(int(backend['provider']))
    if 'host' in backend.keys():
        conn = driver(backend['id'],
                      backend['secret'],
                      False,
                      host=backend['host'],
                      ex_force_auth_url=backend.get('auth_url', None),
                      ex_force_auth_version=backend.get('auth_version', '1.0'),
                      port=80)
    else:
        conn = driver(backend['id'], backend['secret'])
    return conn

@view_config(route_name='home',
             request_method='GET',
             renderer='templates/home.pt')
def home(request):
    """Fill in an object with backend data, taken from config.py"""
    try:
        backend_list = request.environ['beaker.session']['backends']
        session = True
    except:
        backend_list = BACKENDS
        session = False

    backends = []
    for backend in backend_list:
        backends.append({'id'           : backend['id'],
                         'title'        : backend['title'],
                         'provider'     : backend['provider'],
                         'poll_interval': backend['poll_interval'],
                         'status'       : 'off',
                         })

    return {'project': 'mist.io',
            'backends': backends,
            'session': session}

@view_config(route_name='machines', request_method='GET', renderer='json')
def list_machines(request):
    """List machines for a backend"""

    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        machines = conn.list_nodes()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for m in machines:
        # for rackspace get the tags stored in extra.metadata.tags attr, for amazon get extra.tags.tags attr
        tags = m.extra.get('tags', None) or m.extra.get('metadata', None)
        tags = tags and tags.get('tags', None) or []
        machine = {'id'           : m.id,
                  'uuid'          : m.get_uuid(),
                  'name'          : m.name,
                  # both rackspace and amazon have the image in the imageId extra attr,
                  'imageId'       : m.image or m.extra.get('imageId', None),
                  # for rackspace get flavorId extra attr, for amazon the instancetype extra attr
                  'size'          : m.size or m.extra.get('flavorId', None) or m.extra.get('instancetype', None),
                  'state'         : m.state,
                  'private_ips'   : m.private_ips,
                  'public_ips'    : m.public_ips,
                  'tags'          : tags,
                  'extra'         : m.extra,
                  }
        machine.update(get_machine_actions(m, conn))
        ret.append(machine)
    return ret

@view_config(route_name='machines', request_method='POST', renderer='json')
def create_machine(request):
    """Create a new virtual machine on the specified backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        machine_name = request.json_body['name']
        location_id = request.json_body['location']
        image_id = request.json_body['image']
        size_id = request.json_body['size']
    except Exception as e:
        return Response('Invalid payload', 400)

    size = NodeSize(size_id, name='', ram='', disk='', bandwidth='', price='',
                    driver=conn)
    image = NodeImage(image_id, name='', driver=conn)
    location = NodeLocation(location_id, name='', country='', driver=conn)

    if conn.type == Provider.RACKSPACE and len(request.registry.settings['keypairs']):
        # try to deploy node with ssh key installed
        ssh_key = SSHKeyDeployment(request.registry.settings['keypairs'][0][0])
        try:
            conn.deploy_node(name=machine_name,
                             image=image,
                             size=size,
                             location=location,
                             deploy=ssh_key)
            return []
        except:
            log.warn('Failed to deploy node with ssh key. Trying to create simple node')

    try:
        conn.create_node(name=machine_name,
                         image=image,
                         size=size,
                         location=location)
        return []
    except Exception as e:
        return Response('Something went wrong with the creation', 500)

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=start',
             renderer='json')
def start_machine(request):
    """Starts a machine on backends that supported start/stop"""
    #TODO: in which providers could this work?
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    #machine.start()

    return []

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=stop',
             renderer='json')
def stop_machine(request):
    """Stops a machine on backends that supported start/stop"""
    #TODO: in which providers could this work?
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    #machine.stop()

    return []

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=reboot',
             renderer='json')
def reboot_machine(request):
    """Reboots a machine on a certain backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    machine.reboot()
    return []

@view_config(route_name='machine',
             request_method='POST',
             request_param='action=destroy',
             renderer='json')
def destroy_machine(request):
    """Destroys a machine on a certain backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    machine_id = request.matchdict['machine']
    machine = Node(machine_id,
                   name=machine_id,
                   state=0,
                   public_ips=[],
                   private_ips=[],
                   driver=conn)
    machine.destroy()

    return []

@view_config(route_name='metadata', request_method='POST')
def set_metadata(request):
    """Sets metadata for a machine, given the backend and machine id"""
    #TODO: the following are not working
    """Examples
    Openstack:
        conn.ex_set_metadata(machine,
                            {'name': 'ServerX',
                             'description': 'all the money'})
    EC2:
        conn2.ex_create_tags(machine, {'something': 'something_something'})
    """
    ret = []
    done = False
    backends = [b for b in BACKENDS if b['id'] == request.matchdict['backend']]
    if backends:
        backend = backends[0]
        conn = connect(backend)
        machines = conn.list_nodes()
        for machine in machines:
            if machine.id == request.matchdict['machine']:
                try:
                    metadata = request.json_body
                    #get metadata from request
                except:
                    return Response('Not proper format for metadata', 404)
                try:
                    #e.g. Openstack
                    metadata = conn.ex_set_metadata(machine, metadata)
                    done = True
                except:
                    try:
                        #e.g. EC2
                        metadata = conn.ex_create_tags(machine, metadata)
                        done = True
                    except:
                        return Response('Not implemented for this backend', 404)
                break
    if not done:
        return Response('Invalid backend', 404)

    return Response(json.dumps(ret))

@view_config(route_name='images', request_method='GET', renderer='json')
def list_images(request):
    """List images from each backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        try:
            backend_list = request.environ['beaker.session']['backends']
        except:
            backend_list = BACKENDS
        backend_index = int(request.matchdict['backend'])
        backend = backend_list[backend_index]
        if backend['provider'] == Provider.EC2:
            images = conn.list_images(None, BASE_EC2_AMIS.keys())
        else:
            images = conn.list_images()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for image in images:
        ret.append({'id'    : image.id,
                    'extra' : image.extra,
                    'name'  : image.name,
                    })
    return ret

@view_config(route_name='sizes', request_method='GET', renderer='json')
def list_sizes(request):
    """List sizes (aka flavors) from each backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        sizes = conn.list_sizes()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for size in sizes:
        ret.append({'id'        : size.id,
                    'bandwidth' : size.bandwidth,
                    'disk'      : size.disk,
                    'driver'    : size.driver.name,
                    'name'      : size.name,
                    'price'     : size.price,
                    'ram'       : size.ram,
                    })

    return ret

@view_config(route_name='locations', request_method='GET', renderer='json')
def list_locations(request):
    """List locations from each backend"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        locations = conn.list_locations()
    except:
        return Response('Backend unavailable', 503)

    ret = []
    for location in locations:
        ret.append({'id'        : location.id,
                    'name'      : location.name,
                    'country'   : location.country,
                    })

    return ret

@view_config(route_name='image_details', request_method='GET', renderer='json')
def get_image_details(request):
    """get_image_details"""
    try:
        conn = connect(request)
    except:
        return Response('Backend not found', 404)

    try:
        image = conn.get_image(request.params['id'])
    except:
        return Response('Backend unavailable', 503)

    if image is None:
        ret = {}
    else:
        ret = {'id'     : image.id,
               'extra' : image.extra,
               'name'  : image.name,
               }
    return ret

@view_config(route_name='backends', request_method='GET', renderer='json')
def get_backends(request):
    """get_backends"""
    try:
        backend_list = request.environ['beaker.session']['backends']
    except:
        backend_list = BACKENDS

    backends = []
    index = 0
    for backend in backend_list:
        backends.append({'index'        : index,
                         'id'           : backend['id'],
                         'title'        : backend['title'],
                         'provider'     : backend['provider'],
                         'poll_interval': backend['poll_interval'],
                         'status'       : 'online',
                         })
        index = index + 1

    return backends

@view_config(route_name='machine_has_key', request_method='GET', renderer='json')
def machine_has_key(request):
    """Check if the machine has a key pair deployed

    To do that we try to connect to the machine using fabric. The problem is
    that fabric does not support passing the private key as a string, but only
    as a file. To solve this we use a temporary file. After the connection is
    closed we erase it.
    """

    machine_ip = request.params.get('ip', None)
    private_key = request.registry.settings['keypairs'][0][1]

    if not machine_ip or not private_key:
        log.info('IP or private key missing. Skipped checking.')
        return False

    #import pdb; pdb.set_trace()

    env.host_string = machine_ip
    env.user = 'root'
    #env.connection_attempts - defaults to 1
    #env.timeout - e.g. 20 in secs defaults to 10

    (tmp_key, tmp_path) = tempfile.mkstemp()
    key_fd = os.fdopen(tmp_key, 'w+b')
    key_fd.write(private_key)
    key_fd.close()
    env.key_filename = [tmp_path]

    if run('uptime').failed:
        ret = {'has_key': False}
    else:
        ret = {'has_key': True}

    os.remove(tmp_path)

    return ret
<|endoftext|>