<commit_before>from roundup.cgi.actions import Action
from roundup.cgi import exceptions, templating
import cgi
import roundup
import json
import re
from pprint import pprint

def _walkprop(db, node, path):
    n = node
    for p in path.split("."):
        i = n.get(p)

        if p in ['creator']:
            p = 'user'
        n = db.getnode(p, i)
    return n

class GetJSONAction(Action):
    name = 'export'
    permissionType = 'View'

    def _safe_request_lookup(self, name, default=None, cast=None):
        if name in self.request.form:
            val = self.request.form[name].value
            if cast:
                return cast(val)
            return val
        return default

    def handle(self):
        """ Export the specified search query as JSON. """
        # figure the request
        self.request = templating.HTMLRequest(self.client)
        try:
            json_nested = self._safe_request_lookup("json_nested", False, cast=bool)
            itemclass = self._safe_request_lookup("itemclass", self.request.classname, cast=str)
            json_flat_array = self._safe_request_lookup("json_flat_array", False, cast=bool)

            self.columns = {}
            for i, c in enumerate(self.request.columns):
                self.columns[c] = ColumnDescriptor(c, c, i)

            res = self._retrieve_data(itemclass, self.columns.values(), json_nested, self.request.search_text, self.request.filterspec, self.request.sort, self.request.group)
            res = res[0]

            if json_flat_array:
                if len(self.request.columns) > 1:
                    return "ERROR: only one column is allowed when flattening to an array"
                r = []
                for i in res:
                    for k, v in i.items():
                        r.append(v)
                res = r

            if self.client.env['REQUEST_METHOD'] == 'HEAD':
                # all done, return a dummy string
                return 'dummy'
            res = json.dumps(res)
        except:
            self._set_headers("text/html")
            raise
        finally:
            self._set_headers()
        return res

    def _set_headers(self, mimetype='applicaton/json; charset=utf-8'):
        self.client.additional_headers['Content-Type'] = mimetype
        self.client.header()

    def _retrieve_data(self, itemclass, columns, json_nested=False, searchtext=None, filterspec=[], sort=[], group=[]):
        klass = self.db.getclass(itemclass)

        # check if all columns exist on class
        # the exception must be raised before sending header
        props = klass.getprops()
        for cname in columns:
            if not type(cname) is str:
                cname = cname.get_name()
            # TODO ? custom handling of project and device_type, because we need to walk up
            # the tree a little.
            if cname not in props and "." not in cname:
                # TODO raise exceptions.NotFound(.....) does not give message
                # so using SeriousError instead
                self.client.response_code = 404
                raise exceptions.SeriousError(
                    self._('Column "%(column)s" not fo und on %(class)s')
                    % {'column': cgi.escape(cname), 'class': itemclass})

        # full-text search
        # TODO: currently only matches FULL words. This is how roundup implements search.
        matches = None
        if searchtext:
            """matches = self.client.db.indexer.search(
                [w.upper().encode("utf-8", "replace") for w in re.findall(
                    r'(?u)\b\w{2,25}\b',
                    unicode(searchtext, "utf-8", "replace")
                )], klass)

            """
            matches = self.db.indexer.search(
                re.findall(r'\b\w{2,25}\b', searchtext), klass)

        # json_flat is currently broken, always use nested
        if 1 or json_nested:
            obj = self._build_json_nested(klass, columns, klass.filter(matches, filterspec, sort, group))
        else:
            obj = self._build_json_flat(klass, columns, klass.filter(matches, filterspec, sort, group))

        return obj, len(obj), klass.count()

    # this is currently broken!
    def _build_json_flat(self, klass, columns, itemidlist):
        res = []
        for itemid in itemidlist:
            row = {}
            for col in columns:
                name = col.get_name()
                if "." in name:
                    path, d, prop = name.rpartition(".")
                    l = self._walkprop(klass.getnode(itemid), path)
                    row[self.replace_dots_with_dashes and name.replace('.', '-') or name] = l[prop]
                    continue
                row[name] = klass.get(itemid, name)
                if row[name] and isinstance(row[name], roundup.date.Date):
                    row[name] = row[name].formal()
            res.append(row)
        return res

    def _build_tree(self, root, name, node):
        valid_classes = self.db.getclasses()
        tree = root
        fk = None
        for component in name.split('.'):
            if not node:
                # missing foreign key, but still complete the component walk
                # so we get the full transitive tree
                tree[component] = None
                continue
            if component == "id" and fk:
                # we will add a 'fake' id atttribute with the fk of the node itself
                tree[component] = fk
            elif not component in node.keys():
                # TODO: throw a more useful error
                raise Exception("%s is not available on %s" % (component, node))
            fk = node.get(component)
            if not component in tree:
                tree[component] = {}

            if component in ['creator']:
                mapped_component = 'user'
            else:
                mapped_component = component

            if fk is None:
                # there is no foreign key. This is basically a broken reference.
                # for example, if no status is specified on a finding.
                node = None
            elif mapped_component in valid_classes:
                # deal with link
                # TODO: deal with multilinks (created an array in tree?)
                node = self.db.getnode(mapped_component, fk)
            elif component in node:
                # deal with properties:
                tree[component] = node[component]
            else:
                print "WTF!"
                print node
            tree = tree[component]
        return tree

    def _get_project_for_row(self, klass, itemid):
        if klass.classname in ["finding", "antifinding"]:
            path = "version.device_type.project"
        elif klass.classname == "version":
            path = "device_type.project"
        elif klass.classname == "device_type":
            path = "project"
        else:
            return None
        return _walkprop(self.db, klass.getnode(itemid), path)

    def _build_json_nested(self, klass, columns, itemidlist):
        res = []

        for itemid in itemidlist:
            row = {}
            for col in columns:
                pprint(col)
                # deal with transverse
                col_name = col.get_name()
                if "." in col_name:
                    self._build_tree(row, col_name, klass.getnode(itemid))
                    continue

                item = klass.get(itemid, col_name)
                if col_name == "proj_id" and item:
                    project = self._get_project_for_row(klass, itemid)
                    item = "%s%03d" % (project["prefix"], int(item))
                row[col_name] = item
                # special case date:
                if item and isinstance(item, roundup.date.Date):
                    row[col_name] = item.formal()
            res.append(row)
        return res

class ColumnDescriptor:
    def __init__(self, name, db_path, index, sortable=True, searchable=True, search_text=None):
        self.name, self.db_path, self.index, self.sortable, self.searchable, self.search_text = name, db_path, index, sortable, searchable, search_text
        self.is_sorting = False
        self.is_searching = self.searchable and self.search_text
        self.direction = "asc"

    def __repr__(self):
        return "<col_name: %s, db_path: %s, index: %i, sortable: %s, direction: %s>" % (self.name, self.db_path, self.index, self.sortable, self.direction)

    def get_sortattr(self):
        if self.sortable and self.is_sorting:
            return (("-" if self.direction == "desc" else "+"), self.db_path)
        return None

    def get_filterspec(self):
        if self.searchable and self.is_searching:
            return {self.db_path: self.search_text}
        return None

    def get_name(self):
        return self.db_path

class GetJSONForDataTablesAction(GetJSONAction):

    def _parse_client_columns(self):
        columns = self._safe_request_lookup("sColumns", "")
        # TODO: what do we do when there is no client side column?
        for i, col_name in enumerate([c.strip() for c in columns.split(",")]):
            self.columns[i] = ColumnDescriptor(
                col_name,
                self._safe_request_lookup('mDataProp_%d' % i),
                i,
                self._safe_request_lookup('bSortable_%d' % i, True, cast=bool),
                self._safe_request_lookup('bSearchable_%d' % i, True, cast=bool),
                self._safe_request_lookup('sSearch_%d' % i, None, cast=str),
            )
        self.client_columns = [c.strip() for c in columns.split(',')]
        self.num_client_columns = self._safe_request_lookup("iColumns", 0, cast=int)
        num_sort_cols = self._safe_request_lookup("iSortingCols", cast=int)
        if num_sort_cols:
            for i in range(num_sort_cols):
                col_index = self._safe_request_lookup('iSortCol_%d' % i, None, cast=int)
                self.columns[col_index].is_sorting = True
                self.columns[col_index].direction = self._safe_request_lookup('sSortDir_%d' % i, "asc")

    def _num_columns(self):
        return self.num_client_columns

    def _lookup_client_column(self, name):
        return self.client_columns.index(name)

    def handle(self):
        self.columns = {}
        res = ""
        mimetype = "application/json; charset=%s" % self.client.charset
        try:
            self.request = templating.HTMLRequest(self.client)

            # get parameters
            itemclass = self._safe_request_lookup("itemclass", self.request.classname, cast=str)
            json_nested = self._safe_request_lookup("json_nested", False, cast=bool)
            start = self._safe_request_lookup("iDisplayStart", self.request.startwith, cast=int)
            pagesize = self._safe_request_lookup("iDisplayLength", self.request.pagesize, cast=int)
            searchtext = self._safe_request_lookup("sSearch", self.request.search_text)

            self._parse_client_columns()
            pprint(self.columns)

            roundup_sortlist = []
            roundup_filterspec = self.request.filterspec
            for c in self.columns.values():
                sortattr = c.get_sortattr()
                filterspec = c.get_filterspec()
                if sortattr:
                    roundup_sortlist.append(sortattr)
                if filterspec:
                    roundup_filterspec.update(filterspec)

            obj, num_items_selected, total_items = self._retrieve_data(itemclass, self.columns.values(), json_nested, searchtext, roundup_filterspec, roundup_sortlist, self.request.group)

            obj = obj[start:start + pagesize]

            sEcho = self.request.form['sEcho'].value if 'sEcho' in self.request.form else 0
            response = {
                "iTotalDisplayRecords": num_items_selected,
                "iTotalRecords": total_items,
                "sEcho": sEcho,
                "aaData": obj,
            }
            res = json.dumps(response)

        except:
            self._set_headers("text/html")
            raise

        self._set_headers(mimetype)
        return res

def init(instance):
    instance.registerAction('get_json', GetJSONAction)
    instance.registerAction('get_json_for_datatables', GetJSONForDataTablesAction)
<commit_msg>get rid of some special cases<commit_after>from roundup.cgi.actions import Action
from roundup.cgi import exceptions, templating
import cgi
import roundup
import json
import re
from pprint import pprint

def _walkprop(db, node, path):
    n = node
    for p in path.split("."):
        i = n.get(p)

        if p in ['creator']:
            p = 'user'
        n = db.getnode(p, i)
    return n

class GetJSONAction(Action):
    name = 'export'
    permissionType = 'View'

    def _safe_request_lookup(self, name, default=None, cast=None):
        if name in self.request.form:
            val = self.request.form[name].value
            if cast:
                return cast(val)
            return val
        return default

    def handle(self):
        """ Export the specified search query as JSON. """
        # figure the request
        self.request = templating.HTMLRequest(self.client)
        try:
            json_nested = self._safe_request_lookup("json_nested", False, cast=bool)
            itemclass = self._safe_request_lookup("itemclass", self.request.classname, cast=str)
            json_flat_array = self._safe_request_lookup("json_flat_array", False, cast=bool)

            self.columns = {}
            for i, c in enumerate(self.request.columns):
                self.columns[c] = ColumnDescriptor(c, c, i)

            res = self._retrieve_data(itemclass, self.columns.values(), json_nested, self.request.search_text, self.request.filterspec, self.request.sort, self.request.group)
            res = res[0]

            if json_flat_array:
                if len(self.request.columns) > 1:
                    return "ERROR: only one column is allowed when flattening to an array"
                r = []
                for i in res:
                    for k, v in i.items():
                        r.append(v)
                res = r

            if self.client.env['REQUEST_METHOD'] == 'HEAD':
                # all done, return a dummy string
                return 'dummy'
            res = json.dumps(res)
        except:
            self._set_headers("text/html")
            raise
        finally:
            self._set_headers()
        return res

    def _set_headers(self, mimetype='applicaton/json; charset=utf-8'):
        self.client.additional_headers['Content-Type'] = mimetype
        self.client.header()

    def _retrieve_data(self, itemclass, columns, json_nested=False, searchtext=None, filterspec=[], sort=[], group=[]):
        klass = self.db.getclass(itemclass)

        # check if all columns exist on class
        # the exception must be raised before sending header
        props = klass.getprops()
        for cname in columns:
            if not type(cname) is str:
                cname = cname.get_name()
            # TODO ? custom handling of project and device_type, because we need to walk up
            # the tree a little.
            if cname not in props and "." not in cname:
                # TODO raise exceptions.NotFound(.....) does not give message
                # so using SeriousError instead
                self.client.response_code = 404
                raise exceptions.SeriousError(
                    self._('Column "%(column)s" not fo und on %(class)s')
                    % {'column': cgi.escape(cname), 'class': itemclass})

        # full-text search
        # TODO: currently only matches FULL words. This is how roundup implements search.
        matches = None
        if searchtext:
            """matches = self.client.db.indexer.search(
                [w.upper().encode("utf-8", "replace") for w in re.findall(
                    r'(?u)\b\w{2,25}\b',
                    unicode(searchtext, "utf-8", "replace")
                )], klass)

            """
            matches = self.db.indexer.search(
                re.findall(r'\b\w{2,25}\b', searchtext), klass)

        # json_flat is currently broken, always use nested
        if 1 or json_nested:
            obj = self._build_json_nested(klass, columns, klass.filter(matches, filterspec, sort, group))
        else:
            obj = self._build_json_flat(klass, columns, klass.filter(matches, filterspec, sort, group))

        return obj, len(obj), klass.count()

    # this is currently broken!
    def _build_json_flat(self, klass, columns, itemidlist):
        res = []
        for itemid in itemidlist:
            row = {}
            for col in columns:
                name = col.get_name()
                if "." in name:
                    path, d, prop = name.rpartition(".")
                    l = self._walkprop(klass.getnode(itemid), path)
                    row[self.replace_dots_with_dashes and name.replace('.', '-') or name] = l[prop]
                    continue
                row[name] = klass.get(itemid, name)
                if row[name] and isinstance(row[name], roundup.date.Date):
                    row[name] = row[name].formal()
            res.append(row)
        return res

    def _build_tree(self, root, name, node):
        valid_classes = self.db.getclasses()
        tree = root
        fk = None
        for component in name.split('.'):
            if not node:
                # missing foreign key, but still complete the component walk
                # so we get the full transitive tree
                tree[component] = None
                continue
            if component == "id" and fk:
                # we will add a 'fake' id atttribute with the fk of the node itself
                tree[component] = fk
            elif not component in node.keys():
                # TODO: throw a more useful error
                raise Exception("%s is not available on %s" % (component, node))
            fk = node.get(component)
            if not component in tree:
                tree[component] = {}

            if component in ['creator']:
                mapped_component = 'user'
            else:
                mapped_component = component

            if fk is None:
                # there is no foreign key. This is basically a broken reference.
                # for example, if no status is specified on a finding.
                node = None
            elif mapped_component in valid_classes:
                # deal with link
                # TODO: deal with multilinks (created an array in tree?)
                node = self.db.getnode(mapped_component, fk)
            elif component in node:
                # deal with properties:
                tree[component] = node[component]
            else:
                print "WTF!"
                print node
            tree = tree[component]
        return tree

    def _build_json_nested(self, klass, columns, itemidlist):
        res = []

        for itemid in itemidlist:
            row = {}
            for col in columns:
                # deal with transverse
                col_name = col.get_name()
                if "." in col_name:
                    self._build_tree(row, col_name, klass.getnode(itemid))
                    continue

                item = klass.get(itemid, col_name)
                row[col_name] = item
                # special case date:
                if item and isinstance(item, roundup.date.Date):
                    row[col_name] = item.formal()
            res.append(row)
        return res

class ColumnDescriptor:
    def __init__(self, name, db_path, index, sortable=True, searchable=True, search_text=None):
        self.name, self.db_path, self.index, self.sortable, self.searchable, self.search_text = name, db_path, index, sortable, searchable, search_text
        self.is_sorting = False
        self.is_searching = self.searchable and self.search_text
        self.direction = "asc"

    def __repr__(self):
        return "<col_name: %s, db_path: %s, index: %i, sortable: %s, direction: %s>" % (self.name, self.db_path, self.index, self.sortable, self.direction)

    def get_sortattr(self):
        if self.sortable and self.is_sorting:
            return (("-" if self.direction == "desc" else "+"), self.db_path)
        return None

    def get_filterspec(self):
        if self.searchable and self.is_searching:
            return {self.db_path: self.search_text}
        return None

    def get_name(self):
        return self.db_path

class GetJSONForDataTablesAction(GetJSONAction):

    def _parse_client_columns(self):
        columns = self._safe_request_lookup("sColumns", "")
        # TODO: what do we do when there is no client side column?
        for i, col_name in enumerate([c.strip() for c in columns.split(",")]):
            self.columns[i] = ColumnDescriptor(
                col_name,
                self._safe_request_lookup('mDataProp_%d' % i),
                i,
                self._safe_request_lookup('bSortable_%d' % i, True, cast=bool),
                self._safe_request_lookup('bSearchable_%d' % i, True, cast=bool),
                self._safe_request_lookup('sSearch_%d' % i, None, cast=str),
            )
        self.client_columns = [c.strip() for c in columns.split(',')]
        self.num_client_columns = self._safe_request_lookup("iColumns", 0, cast=int)
        num_sort_cols = self._safe_request_lookup("iSortingCols", cast=int)
        if num_sort_cols:
            for i in range(num_sort_cols):
                col_index = self._safe_request_lookup('iSortCol_%d' % i, None, cast=int)
                self.columns[col_index].is_sorting = True
                self.columns[col_index].direction = self._safe_request_lookup('sSortDir_%d' % i, "asc")

    def _num_columns(self):
        return self.num_client_columns

    def _lookup_client_column(self, name):
        return self.client_columns.index(name)

    def handle(self):
        self.columns = {}
        res = ""
        mimetype = "application/json; charset=%s" % self.client.charset
        try:
            self.request = templating.HTMLRequest(self.client)

            # get parameters
            itemclass = self._safe_request_lookup("itemclass", self.request.classname, cast=str)
            json_nested = self._safe_request_lookup("json_nested", False, cast=bool)
            start = self._safe_request_lookup("iDisplayStart", self.request.startwith, cast=int)
            pagesize = self._safe_request_lookup("iDisplayLength", self.request.pagesize, cast=int)
            searchtext = self._safe_request_lookup("sSearch", self.request.search_text)

            self._parse_client_columns()

            roundup_sortlist = []
            roundup_filterspec = self.request.filterspec
            for c in self.columns.values():
                sortattr = c.get_sortattr()
                filterspec = c.get_filterspec()
                if sortattr:
                    roundup_sortlist.append(sortattr)
                if filterspec:
                    roundup_filterspec.update(filterspec)

            obj, num_items_selected, total_items = self._retrieve_data(itemclass, self.columns.values(), json_nested, searchtext, roundup_filterspec, roundup_sortlist, self.request.group)

            obj = obj[start:start + pagesize]

            sEcho = self.request.form['sEcho'].value if 'sEcho' in self.request.form else 0
            response = {
                "iTotalDisplayRecords": num_items_selected,
                "iTotalRecords": total_items,
                "sEcho": sEcho,
                "aaData": obj,
            }
            res = json.dumps(response)

        except:
            self._set_headers("text/html")
            raise

        self._set_headers(mimetype)
        return res

def init(instance):
    instance.registerAction('get_json', GetJSONAction)
    instance.registerAction('get_json_for_datatables', GetJSONForDataTablesAction)
<|endoftext|>


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

import netmap       # our module
import time         # time measurements
import select       # poll()
import argparse     # program argument parsing
import multiprocessing    # thread management
import re

# import scapy suppressing the initial WARNING message
import logging
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
from scapy.all import Ether, IP, UDP  # packet forgery

def help_quit(parser):
    print("")
    parser.print_help()
    quit()

def build_packet(args, parser):
    src = args.src.split(':')
    dst = args.dst.split(':')

    # create the payload
    base = "Hello from Python"
    header_len = 14 + 20 + 8
    data = base * ((args.length-header_len)/len(base) + 1)
    data = data[0:args.length-header_len]

    scap = Ether(src = args.srcmac, dst = args.dstmac)
    scap = scap / IP(src = src[0], dst = dst[0])
    scap = scap / UDP(sport = int(src[1]), dport = int(dst[1]))
    scap = scap / data

    try:
        # checksum is computed when calling str(scap), e.g. when the packet is
        # assembled
        ret = str(scap)
    except:
        print("Packet parameters are invalid\n")
        help_quit(parser)

    if args.dump:
        scap.show2()

    return ret

def transmit(idx, suffix, args, parser, queue):
    # use nm_open() to open the netmap device and register an interface
    # using an extended interface name
    nmd = netmap.NetmapDesc(args.interface + suffix)
    time.sleep(args.wait_link)

    # build the packet that will be transmitted
    pkt = build_packet(args, parser)

    # fill in the netmap slots and netmap buffers for tx ring 0
    txr = nmd.transmit_rings[idx]
    num_slots = txr.num_slots
    for i in range(num_slots):
        txr.slots[i].buf[0:len(pkt)] = pkt
        txr.slots[i].len = len(pkt)

    # transmit at maximum speed until Ctr-C is pressed
    cnt = 0         # packet counter
    batch = args.batch
    poller = select.poll()
    poller.register(nmd.getfd(), select.POLLOUT)
    t_start = time.time()
    try:
        cur = txr.cur
        while 1:
            ready_list = poller.poll(2)
            if len(ready_list) == 0:
                print("Timeout occurred")
                break;
            n = txr.tail - cur  # avail
            if n < 0:
                n += num_slots
            if n > batch:
                n = batch
            cur += n
            if cur >= num_slots:
                cur -= num_slots
            txr.cur = txr.head = cur # lazy update txr.cur and txr.head
            nmd.txsync()
            cnt += n
    except KeyboardInterrupt:
        # report the result to the main process
        queue.put([cnt, time.time() - t_start])
        pass

def receive(idx, suffix, args, parser, queue):
    # use nm_open() to open the netmap device and register an interface
    # using an extended interface name
    nmd = netmap.NetmapDesc(args.interface + suffix)
    time.sleep(args.wait_link)

    # select the right ring
    rxr = nmd.receive_rings[idx]
    num_slots = rxr.num_slots

    cnt = 0         # packet counter
    poller = select.poll()
    poller.register(nmd.getfd(), select.POLLIN)

    # wait for the first packet
    try:
        poller.poll()
    except KeyboardInterrupt:
        # report the result to the main process
        queue.put([cnt, None])
        return

    # receive (throwing away everything) until Ctr-C is pressed
    t_start = time.time()
    try:
        cur = rxr.cur
        while 1:
            ready_list = poller.poll()
            if len(ready_list) == 0:
                print("Timeout occurred")
                break;
            n = rxr.tail - cur  # avail
            if n < 0:
                n += num_slots
            cur += n
            if cur >= num_slots:
                cur -= num_slots
            rxr.cur = rxr.head = cur # lazy update rxr.cur and rxr.head
            cnt += n
    except KeyboardInterrupt:
        # report the result to the main process
        queue.put([cnt, time.time() - t_start])
        pass

# How many netmap ring couples has 'ifname'?
def netmap_max_rings(ifname):
    if ifname.startswith('netmap:'):
        ifname = ifname[7:]

    nm = netmap.Netmap()
    nm.open()
    nm.if_name = ifname
    nm.getinfo()

    return nm.tx_rings

# extract the (nr_ringid, nr_flags) specified by the extended
# interface name (nm_open() ifname)
def netmap_get_ringid(ifname):
    if ifname.startswith('netmap:'):
        ifname = ifname[7:]

    nm = netmap.Netmap()
    nm.open()
    nm.if_name = ifname
    nm.getinfo()

    return nm.ringid, nm.flags

def netmap_remove_ifname_suffix(ifname_ext):
    m = re.match(r'\w+:\w+', ifname_ext)
    if m == None:
        return None

    return m.group(0)

############################## MAIN ###########################

if __name__ == '__main__':

    # functions implemented by this program
    handler = dict();
    handler['tx'] = transmit
    handler['rx'] = receive

    # program arguments
    parser = argparse.ArgumentParser(description = 'Send and receive packet using the netmap API')
    parser.add_argument('-i', '--interface', help = 'the interface to register with netmap; '
                        'can be in the form netmap:<OSNAME>[<EXT>] or <VALENAME>[<EXT>], where '
                        'OSNAME is the O.S. name for a network interface (e.g. "eth0"), '
                        '<VALENAME> is a valid VALE port name (e.g. "vale18:2") and <EXT> is an '
                        'optional extension suffix, specified using the nm_open() syntax '
                        '(e.g. "^", "-5", "{44", ...)',
                    required = True)
    parser.add_argument('-f', '--function', help = 'the function to perform',
                    choices = ['tx', 'rx'], default = 'rx')
    parser.add_argument('-b', '--batchsize', help = 'number of packets to send with each TXSYNC '
                    'operation', type=int, default = 512, dest = 'batch')
    parser.add_argument('-l', '--length', help = 'lenght of the ethernet frame sent',
                    type = int, default = 60)
    parser.add_argument('-D', '--dstmac', help = 'destination MAC of tx packets',
                    default = 'ff:ff:ff:ff:ff:ff')
    parser.add_argument('-S', '--srcmac', help = 'source MAC of tx packets',
                    default = '00:00:00:00:00:00')
    parser.add_argument('-d', '--dst', help = 'destination IP address and UDP port of tx packets',
                    default = '10.0.0.2:54322', metavar = 'IP:PORT')
    parser.add_argument('-s', '--src', help = 'source IP address and UDP port of tx packets',
                    default = '10.0.0.1:54321', metavar = 'IP:PORT')
    parser.add_argument('-w', '--wait-link', help = 'time to wait for the link before starting '
                    'transmit/receive operations (in seconds)', type = int, default = 1)
    parser.add_argument('-X', '--dump', help = 'dump the packet', action = 'store_true')
    parser.add_argument('-p', '--threads', help = 'number of threads to used for tx/rx '
                    'operations', type = int, default = 1)
    # parse the input
    args = parser.parse_args()
    # print args

    # bound checking
    if args.length < 60:
        print('Invalid packet length\n')
        help_quit(parser)

    if args.threads < 1:
        print('Invalid number of threads\n')
        help_quit(parser)

    try:
        # compute 'ifname' removing the suffix from the extended name
        # specified by the user
        ifname = netmap_remove_ifname_suffix(args.interface)
        if ifname == None:
            print('Invalid ifname "%s"' % (args.interface, ))
            help_quit(parser)

        # compute 'max_couples', which is the number of tx/rx rings couples to be registered
        # according to 'args.interface'
        nr_ringid, nr_flags = netmap_get_ringid(args.interface)
        if nr_flags in [netmap.RegAllNic, netmap.RegNicSw]:
            # ask netmap for the number of available couples
            max_couples = netmap_max_rings(args.interface)
            suffix_required = True
            ringid_offset = 0
        else:
            # all the others netmap.Reg* specifies just one couple of rings
            max_couples = 1
            suffix_required = False
            ringid_offset = nr_ringid
        if args.threads > max_couples:
            print('You cannot use more than %s (tx,rx) rings couples with "%s"' % (max_couples, args.interface))
            help_quit(parser)
    except netmap.error as e:
        print(e)
        quit()

    jobs = []    # array of worker processes
    queues = []  # array of queues for IPC
    for i in range(args.threads):
        queue = multiprocessing.Queue()
        queues.append(queue)

        # 'i_off' contains the ring idx on which the process below will operate
        i_off = i + ringid_offset
        # it may also be necessary to add an extension suffix to the interface
        # name specified by the user
        if suffix_required:
            suffix = '-' + str(i_off)
        else:
            suffix = ''

        # create a new process that will execute the user-selected handler function,
        # with the arguments specified by the 'args' tuple
        job = multiprocessing.Process(name = 'worker-' + str(i),
                                        target = handler[args.function],
                                        args = (i_off, suffix, args, parser, queue))
        job.deamon = True   # ensure work termination
        jobs.append(job)

    # start all the workers
    for i in range(len(jobs)):
        jobs[i].start()

    # Wait for the user pressing Ctrl-C
    try:
        while 1:
            time.sleep(1000)
    except KeyboardInterrupt:
        pass

    # collect and print the result returned by the workers
    tot_rate = 0.0
    for i in range(len(jobs)):
        result = queues[i].get()
        jobs[i].join()
        delta = result[1]
        cnt = result[0]
        if delta == None:
            rate = None
        else:
            rate = 0.001 * cnt / delta
            tot_rate += rate
        print('[%d] Packets processed: %s, Avg rate %s Kpps' % (i, cnt, rate))
    print('Total rate: %s' % (tot_rate, ))
<commit_msg>python2: pktman.py: fix typo<commit_after>#!/usr/bin/env python

import netmap       # our module
import time         # time measurements
import select       # poll()
import argparse     # program argument parsing
import multiprocessing    # thread management
import re

# import scapy suppressing the initial WARNING message
import logging
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
from scapy.all import Ether, IP, UDP  # packet forgery

def help_quit(parser):
    print("")
    parser.print_help()
    quit()

def build_packet(args, parser):
    src = args.src.split(':')
    dst = args.dst.split(':')

    # create the payload
    base = "Hello from Python"
    header_len = 14 + 20 + 8
    data = base * ((args.length-header_len)/len(base) + 1)
    data = data[0:args.length-header_len]

    scap = Ether(src = args.srcmac, dst = args.dstmac)
    scap = scap / IP(src = src[0], dst = dst[0])
    scap = scap / UDP(sport = int(src[1]), dport = int(dst[1]))
    scap = scap / data

    try:
        # checksum is computed when calling str(scap), e.g. when the packet is
        # assembled
        ret = str(scap)
    except:
        print("Packet parameters are invalid\n")
        help_quit(parser)

    if args.dump:
        scap.show2()

    return ret

def transmit(idx, suffix, args, parser, queue):
    # use nm_open() to open the netmap device and register an interface
    # using an extended interface name
    nmd = netmap.NetmapDesc(args.interface + suffix)
    time.sleep(args.wait_link)

    # build the packet that will be transmitted
    pkt = build_packet(args, parser)

    # fill in the netmap slots and netmap buffers for tx ring 0
    txr = nmd.transmit_rings[idx]
    num_slots = txr.num_slots
    for i in range(num_slots):
        txr.slots[i].buf[0:len(pkt)] = pkt
        txr.slots[i].len = len(pkt)

    # transmit at maximum speed until Ctr-C is pressed
    cnt = 0         # packet counter
    batch = args.batch
    poller = select.poll()
    poller.register(nmd.getfd(), select.POLLOUT)
    t_start = time.time()
    try:
        cur = txr.cur
        while 1:
            ready_list = poller.poll(2)
            if len(ready_list) == 0:
                print("Timeout occurred")
                break;
            n = txr.tail - cur  # avail
            if n < 0:
                n += num_slots
            if n > batch:
                n = batch
            cur += n
            if cur >= num_slots:
                cur -= num_slots
            txr.cur = txr.head = cur # lazy update txr.cur and txr.head
            nmd.txsync()
            cnt += n
    except KeyboardInterrupt:
        # report the result to the main process
        queue.put([cnt, time.time() - t_start])
        pass

def receive(idx, suffix, args, parser, queue):
    # use nm_open() to open the netmap device and register an interface
    # using an extended interface name
    nmd = netmap.NetmapDesc(args.interface + suffix)
    time.sleep(args.wait_link)

    # select the right ring
    rxr = nmd.receive_rings[idx]
    num_slots = rxr.num_slots

    cnt = 0         # packet counter
    poller = select.poll()
    poller.register(nmd.getfd(), select.POLLIN)

    # wait for the first packet
    try:
        poller.poll()
    except KeyboardInterrupt:
        # report the result to the main process
        queue.put([cnt, None])
        return

    # receive (throwing away everything) until Ctr-C is pressed
    t_start = time.time()
    try:
        cur = rxr.cur
        while 1:
            ready_list = poller.poll()
            if len(ready_list) == 0:
                print("Timeout occurred")
                break;
            n = rxr.tail - cur  # avail
            if n < 0:
                n += num_slots
            cur += n
            if cur >= num_slots:
                cur -= num_slots
            rxr.cur = rxr.head = cur # lazy update rxr.cur and rxr.head
            cnt += n
    except KeyboardInterrupt:
        # report the result to the main process
        queue.put([cnt, time.time() - t_start])
        pass

# How many netmap ring couples has 'ifname'?
def netmap_max_rings(ifname):
    if ifname.startswith('netmap:'):
        ifname = ifname[7:]

    nm = netmap.Netmap()
    nm.open()
    nm.if_name = ifname
    nm.getinfo()

    return nm.tx_rings

# extract the (nr_ringid, nr_flags) specified by the extended
# interface name (nm_open() ifname)
def netmap_get_ringid(ifname):
    if ifname.startswith('netmap:'):
        ifname = ifname[7:]

    nm = netmap.Netmap()
    nm.open()
    nm.if_name = ifname
    nm.getinfo()

    return nm.ringid, nm.flags

def netmap_remove_ifname_suffix(ifname_ext):
    m = re.match(r'\w+:\w+', ifname_ext)
    if m == None:
        return None

    return m.group(0)

############################## MAIN ###########################

if __name__ == '__main__':

    # functions implemented by this program
    handler = dict();
    handler['tx'] = transmit
    handler['rx'] = receive

    # program arguments
    parser = argparse.ArgumentParser(description = 'Send or receive packets using the netmap API')
    parser.add_argument('-i', '--interface', help = 'the interface to register with netmap; '
                        'can be in the form netmap:<OSNAME>[<EXT>] or <VALENAME>[<EXT>], where '
                        'OSNAME is the O.S. name for a network interface (e.g. "eth0"), '
                        '<VALENAME> is a valid VALE port name (e.g. "vale18:2") and <EXT> is an '
                        'optional extension suffix, specified using the nm_open() syntax '
                        '(e.g. "^", "-5", "{44", ...)',
                    required = True)
    parser.add_argument('-f', '--function', help = 'the function to perform',
                    choices = ['tx', 'rx'], default = 'rx')
    parser.add_argument('-b', '--batchsize', help = 'number of packets to send with each TXSYNC '
                    'operation', type=int, default = 512, dest = 'batch')
    parser.add_argument('-l', '--length', help = 'lenght of the ethernet frame sent',
                    type = int, default = 60)
    parser.add_argument('-D', '--dstmac', help = 'destination MAC of tx packets',
                    default = 'ff:ff:ff:ff:ff:ff')
    parser.add_argument('-S', '--srcmac', help = 'source MAC of tx packets',
                    default = '00:00:00:00:00:00')
    parser.add_argument('-d', '--dst', help = 'destination IP address and UDP port of tx packets',
                    default = '10.0.0.2:54322', metavar = 'IP:PORT')
    parser.add_argument('-s', '--src', help = 'source IP address and UDP port of tx packets',
                    default = '10.0.0.1:54321', metavar = 'IP:PORT')
    parser.add_argument('-w', '--wait-link', help = 'time to wait for the link before starting '
                    'transmit/receive operations (in seconds)', type = int, default = 1)
    parser.add_argument('-X', '--dump', help = 'dump the packet', action = 'store_true')
    parser.add_argument('-p', '--threads', help = 'number of threads to used for tx/rx '
                    'operations', type = int, default = 1)
    # parse the input
    args = parser.parse_args()
    # print args

    # bound checking
    if args.length < 60:
        print('Invalid packet length\n')
        help_quit(parser)

    if args.threads < 1:
        print('Invalid number of threads\n')
        help_quit(parser)

    try:
        # compute 'ifname' removing the suffix from the extended name
        # specified by the user
        ifname = netmap_remove_ifname_suffix(args.interface)
        if ifname == None:
            print('Invalid ifname "%s"' % (args.interface, ))
            help_quit(parser)

        # compute 'max_couples', which is the number of tx/rx rings couples to be registered
        # according to 'args.interface'
        nr_ringid, nr_flags = netmap_get_ringid(args.interface)
        if nr_flags in [netmap.RegAllNic, netmap.RegNicSw]:
            # ask netmap for the number of available couples
            max_couples = netmap_max_rings(args.interface)
            suffix_required = True
            ringid_offset = 0
        else:
            # all the others netmap.Reg* specifies just one couple of rings
            max_couples = 1
            suffix_required = False
            ringid_offset = nr_ringid
        if args.threads > max_couples:
            print('You cannot use more than %s (tx,rx) rings couples with "%s"' % (max_couples, args.interface))
            help_quit(parser)
    except netmap.error as e:
        print(e)
        quit()

    jobs = []    # array of worker processes
    queues = []  # array of queues for IPC
    for i in range(args.threads):
        queue = multiprocessing.Queue()
        queues.append(queue)

        # 'i_off' contains the ring idx on which the process below will operate
        i_off = i + ringid_offset
        # it may also be necessary to add an extension suffix to the interface
        # name specified by the user
        if suffix_required:
            suffix = '-' + str(i_off)
        else:
            suffix = ''

        # create a new process that will execute the user-selected handler function,
        # with the arguments specified by the 'args' tuple
        job = multiprocessing.Process(name = 'worker-' + str(i),
                                        target = handler[args.function],
                                        args = (i_off, suffix, args, parser, queue))
        job.deamon = True   # ensure work termination
        jobs.append(job)

    # start all the workers
    for i in range(len(jobs)):
        jobs[i].start()

    # Wait for the user pressing Ctrl-C
    try:
        while 1:
            time.sleep(1000)
    except KeyboardInterrupt:
        pass

    # collect and print the result returned by the workers
    tot_rate = 0.0
    for i in range(len(jobs)):
        result = queues[i].get()
        jobs[i].join()
        delta = result[1]
        cnt = result[0]
        if delta == None:
            rate = None
        else:
            rate = 0.001 * cnt / delta
            tot_rate += rate
        print('[%d] Packets processed: %s, Avg rate %s Kpps' % (i, cnt, rate))
    print('Total rate: %s' % (tot_rate, ))
<|endoftext|>


<commit_before>import codecs
import warnings
from scrapy.contrib.exporter import BaseItemExporter
import re
import os

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: http://doc.scrapy.org/en/latest/topics/item-pipeline.html
uni = lambda x: unicode(x)

def iterlines(string):
    return iter(string.splitlines())

def extract_filename(string):
    value = re.match(u'(.+?)\.(?:dll|h|lib)', string)
    if value is not None:
        value = uni(value.group(1))
    return (value or '')

def uopen(file, mode='r', buffering=-1, encoding=None,
          errors=None, newline=None, closefd=True, opener=None):
    if newline is not None:
        warnings.warn('newline is not supported in py2')
    if not closefd:
        warnings.warn('closefd is not supported in py2')
    if opener is not None:
        warnings.warn('opener is not supported in py2')
    return codecs.open(filename=file, mode=mode, encoding=encoding,
                       errors=errors, buffering=buffering)

def serialize_metadata(metadata):
    pre_result = ['// Metadata:']
    for key, value in metadata.items():
        pre_result.append(
            u'//   {0}: {1}'.format(
                uni(key.capitalize()),
                uni(value)
            )
        )
    return pre_result

enum_entry_matcher = u'(\w+)\s+((?:0x[0-9a-fA-F]+)|\d+)\s*(.*)'
enum_start_line = u'typedef enum EnumType{0} {{ // {1}'.format
enum_value_line = u'  {0} = {1},'.format
enum_end_line = u'};'

def serialize_tables(tables):
    if len(tables) == 0:
        return ''
    other_entries = [u'// Tables:']
    enum_entries = []
    enum_count = 1
    for header, table in tables:
        # Case 1 - It's an enum
        entry_lines = [enum_start_line(enum_count, uni(header))]
        for entry in table:
            match = re.match(enum_entry_matcher, entry)
            if match is None:
                break
            line = enum_value_line(
                uni(match.group(1)),
                uni(match.group(2))
            )
            if match.group(3):
                line += u' // {0}'.format(uni(match.group(3)))
            entry_lines.append(line)
        else:
            enum_entries.extend(entry_lines)
            enum_entries.append(enum_end_line)
            enum_count += 1
            continue

        # Case 2 - Not an enum
        other_entries.append(u'//   {0}'.format(uni(header)))
        for entry in table:
            other_entries.append(u"//     '{0}',".format(uni(entry)))

    if len(other_entries) == 1:
        return enum_entries
    other_entries.extend(enum_entries)
    return other_entries

class ApiExporter(BaseItemExporter):

    def __init__(self, file_handle, **kwargs):
        self._configure(kwargs, dont_fail=True)
        self.file = file_handle

    def export_item(self, item):
        pre_result = []
        pre_result.extend(serialize_tables(item['tables']))
        pre_result.append(item['code'])
        pre_result.append(u'//')
        pre_result.append(u'// URL: {0}'.format(item['url']))
        pre_result.extend(serialize_metadata(item['metadata']))
        pre_result.append("\n\n")
        contents = u'\n'.join(pre_result).replace(u'\x00', u'').encode('utf-8')

        old_position = file.tell()
        self.file.write(contents)
        self.file.flush()
        os.fsync(self.file.fileno()),
        assert(old_position + len(contents) == file.tell())

class ApiExporterDict(dict):

    def __missing__(self, key):
        value = ApiExporter(open(key, 'w'))
        self[key] = value
        return value

class ApiExportPipeline(object):

    def __init__(self):
        self.links_log = open('crawled_links.txt', 'w')
        self.exporters = ApiExporterDict()
        self.spiders = 0

    def open_spider(self, spider):
        self.spiders += 1

    def close_spider(self, spider):
        self.spiders -= 1
        if self.spiders == 0:
            for exporter in self.exporters.items():
                exporter.file.close()

    def process_item(self, item, spider):
        is_empty = (
            len(item['code']) == 0 and
            len(item['metadata']) == 0 and
            len(item['tables']) == 0
        )

        if is_empty:
            log_line = "Empty URL: '{0}'\n".format(item['url'])
        else:
            log_line = "Non-Empty URL: '{0}'\n".format(item['url'])
            get_metadata = item['metadata'].get
            file_name = (
                get_metadata('header') or
                get_metadata('dll') or
                get_metadata('library') or
                ''
            )
            file_name = extract_filename(file_name) or 'other'
            file_name += '.txt'

            exporter = self.exporters[file_name]
            exporter.export_item(item)

        self.links_log.write(log_line.replace('\x00', '').encode("utf-8"))
        return item
<commit_msg>Switch to new file handling scheme<commit_after>import codecs
import warnings
from scrapy.contrib.exporter import BaseItemExporter
import re
import os

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: http://doc.scrapy.org/en/latest/topics/item-pipeline.html
uni = lambda x: unicode(x)

def iterlines(string):
    return iter(string.splitlines())

def extract_filename(string):
    value = re.match(u'(.+?)\.(?:dll|h|lib)', string)
    if value is not None:
        value = uni(value.group(1))
    return (value or '')

def uopen(file, mode='r', buffering=-1, encoding=None,
          errors=None, newline=None, closefd=True, opener=None):
    if newline is not None:
        warnings.warn('newline is not supported in py2')
    if not closefd:
        warnings.warn('closefd is not supported in py2')
    if opener is not None:
        warnings.warn('opener is not supported in py2')
    return codecs.open(filename=file, mode=mode, encoding=encoding,
                       errors=errors, buffering=buffering)

def serialize_metadata(metadata):
    pre_result = ['// Metadata:']
    for key, value in metadata.items():
        pre_result.append(
            u'//   {0}: {1}'.format(
                uni(key.capitalize()),
                uni(value)
            )
        )
    return pre_result

enum_entry_matcher = u'(\w+)\s+((?:0x[0-9a-fA-F]+)|\d+)\s*(.*)'
enum_start_line = u'typedef enum EnumType{0} {{ // {1}'.format
enum_value_line = u'  {0} = {1},'.format
enum_end_line = u'};'

def serialize_tables(tables):
    if len(tables) == 0:
        return ''
    other_entries = [u'// Tables:']
    enum_entries = []
    enum_count = 1
    for header, table in tables:
        # Case 1 - It's an enum
        entry_lines = [enum_start_line(enum_count, uni(header))]
        for entry in table:
            match = re.match(enum_entry_matcher, entry)
            if match is None:
                break
            line = enum_value_line(
                uni(match.group(1)),
                uni(match.group(2))
            )
            if match.group(3):
                line += u' // {0}'.format(uni(match.group(3)))
            entry_lines.append(line)
        else:
            enum_entries.extend(entry_lines)
            enum_entries.append(enum_end_line)
            enum_count += 1
            continue

        # Case 2 - Not an enum
        other_entries.append(u'//   {0}'.format(uni(header)))
        for entry in table:
            other_entries.append(u"//     '{0}',".format(uni(entry)))

    if len(other_entries) == 1:
        return enum_entries
    other_entries.extend(enum_entries)
    return other_entries

class ApiExporter(BaseItemExporter):

    def __init__(self, file_name, **kwargs):
        self._configure(kwargs, dont_fail=True)
        self.file = open(file_name, 'w+')

    def close():
        self.file.close()

    def export_item(self, item):
        pre_result = []
        pre_result.extend(serialize_tables(item['tables']))
        pre_result.append(item['code'])
        pre_result.append(u'//')
        pre_result.append(u'// URL: {0}'.format(item['url']))
        pre_result.extend(serialize_metadata(item['metadata']))
        pre_result.append("\n\n")
        contents = u'\n'.join(pre_result).replace(u'\x00', u'').encode('utf-8')

        old_position = file.tell()
        self.file.write(contents)
        self.file.flush()
        os.fsync(self.file.fileno()),
        assert(old_position + len(contents) == file.tell())

class ApiExporterDict(dict):

    def __missing__(self, key):
        value = ApiExporter(open(key, 'w'))
        self[key] = value
        return value

class ApiExportPipeline(object):

    def __init__(self):
        self.links_log = open('crawled_links.txt', 'w')
        #self.exporters = ApiExporterDict()
        self.current_file_name = ''
        self.current_file = None
        self.spiders = 0

    def open_spider(self, spider):
        self.spiders += 1

    def close_spider(self, spider):
        self.spiders -= 1
        if self.spiders == 0:
            for exporter in self.exporters.items():
                exporter.file.close()
            self.exporters = ApiExporterDict()

    def process_item(self, item, spider):
        is_empty = (
            len(item['code']) == 0 and
            len(item['metadata']) == 0 and
            len(item['tables']) == 0
        )

        if is_empty:
            log_line = "Empty URL: '{0}'\n".format(item['url'])
        else:
            log_line = "Non-Empty URL: '{0}'\n".format(item['url'])
            get_metadata = item['metadata'].get
            file_name = (
                get_metadata('header') or
                get_metadata('dll') or
                get_metadata('library') or
                ''
            )
            file_name = extract_filename(file_name) or 'other'
            file_name += '.txt'

            if file_name != self.current_file_name:
                if self.exporter is not None:
                    self.exporter.close()
                self.exporter = ApiExporter(file_name)
            self.exporter.export_item(item)

        self.links_log.write(log_line.replace('\x00', '').encode("utf-8"))
        return item
<|endoftext|>


<commit_before>import json

from celery import task
import github

from django.db import transaction
from ide.models.user import User
from ide.models.project import Project
from ide.models.files import SourceFile
from ide.utils.sdk import dict_to_pretty_json
from ide.utils import generate_half_uuid
from utils.keen_helper import send_keen_event

@task(acks_late=True)
def import_gist(user_id, gist_id):
    user = User.objects.get(pk=user_id)
    g = github.Github()

    try:
        gist = g.get_gist(gist_id)
    except github.UnknownObjectException:
        send_keen_event('cloudpebble', 'cloudpebble_gist_not_found', user=user, data={'data': {'gist_id': gist_id}})
        raise Exception("Couldn't find gist to import.")

    files = gist.files
    default_name = gist.description or 'Sample project'

    is_native = True

    if 'appinfo.json' in files:
        settings = json.loads(files['appinfo.json'].content)
        if len(files) == 2 and 'simply.js' in files:
            is_native = False
    else:
        settings = {}
        if len(files) == 1 and 'simply.js' in files:
            is_native = False

    project_settings = {
        'name': settings.get('longName', default_name),
        'owner': user,
        'app_uuid':  generate_half_uuid(),
        'app_short_name': settings.get('shortName', default_name),
        'app_long_name': settings.get('longName', default_name),
        'app_company_name': settings.get('companyName', user.username),
        'app_version_code': 1,
        'app_version_label': settings.get('versionLabel', '1.0'),
        'app_is_watchface': settings.get('watchapp', {}).get('watchface', False),
        'app_capabilities': ','.join(settings.get('capabilities', [])),
        'app_keys': dict_to_pretty_json(settings.get('appKeys', {})),
        'project_type': 'native' if is_native else 'simplyjs'
    }

    with transaction.commit_on_success():
        project = Project.objects.create(**project_settings)

        if is_native:
            for filename in gist.files:
                if filename.endswith('.c') or filename.endswith('.h') or filename == 'pebble-js-app.js':
                    # Because gists can't have subdirectories.
                    if filename == 'pebble-js-app.js':
                        cp_filename = 'js/pebble-js-app.js'
                    else:
                        cp_filename = filename
                    source_file = SourceFile.objects.create(project=project, file_name=cp_filename)
                    source_file.save_file(gist.files[filename].content)
        else:
            source_file = SourceFile.objects.create(project=project, file_name='app.js')
            source_file.save_file(gist.files['simply.js'].content)

    send_keen_event('cloudpebble', 'cloudpebble_gist_import', project=project, data={'data': {'gist_id': gist_id}})
    return project.id
<commit_msg>PBL-6930: Support importing resources from gists.<commit_after>import json

from celery import task
import github

from django.db import transaction
from ide.models.user import User
from ide.models.project import Project
from ide.models.files import SourceFile, ResourceFile, ResourceIdentifier
from ide.utils.sdk import dict_to_pretty_json
from ide.utils import generate_half_uuid
from utils.keen_helper import send_keen_event
import urllib2

@task(acks_late=True)
def import_gist(user_id, gist_id):
    user = User.objects.get(pk=user_id)
    g = github.Github()

    try:
        gist = g.get_gist(gist_id)
    except github.UnknownObjectException:
        send_keen_event('cloudpebble', 'cloudpebble_gist_not_found', user=user, data={'data': {'gist_id': gist_id}})
        raise Exception("Couldn't find gist to import.")

    files = gist.files
    default_name = gist.description or 'Sample project'

    is_native = True

    if 'appinfo.json' in files:
        settings = json.loads(files['appinfo.json'].content)
        if len(files) == 2 and 'simply.js' in files:
            is_native = False
    else:
        settings = {}
        if len(files) == 1 and 'simply.js' in files:
            is_native = False

    project_settings = {
        'name': settings.get('longName', default_name),
        'owner': user,
        'app_uuid':  generate_half_uuid(),
        'app_short_name': settings.get('shortName', default_name),
        'app_long_name': settings.get('longName', default_name),
        'app_company_name': settings.get('companyName', user.username),
        'app_version_code': 1,
        'app_version_label': settings.get('versionLabel', '1.0'),
        'app_is_watchface': settings.get('watchapp', {}).get('watchface', False),
        'app_capabilities': ','.join(settings.get('capabilities', [])),
        'app_keys': dict_to_pretty_json(settings.get('appKeys', {})),
        'project_type': 'native' if is_native else 'simplyjs'
    }

    with transaction.commit_on_success():
        project = Project.objects.create(**project_settings)

        if is_native:
            for filename in gist.files:
                if filename.endswith('.c') or filename.endswith('.h') or filename.endswith('.js'):
                    # Because gists can't have subdirectories.
                    if filename == 'pebble-js-app.js':
                        cp_filename = 'js/pebble-js-app.js'
                    else:
                        cp_filename = filename
                    source_file = SourceFile.objects.create(project=project, file_name=cp_filename)
                    source_file.save_file(gist.files[filename].content)

            media = settings.get('resources', {}).get('media', [])
            resources = {}
            for resource in media:
                kind = resource['type']
                def_name = resource['name']
                filename = resource['file']
                regex = resource.get('characterRegex', None)
                tracking = resource.get('trackingAdjust', None)
                is_menu_icon = resource.get('menuIcon', False)
                if filename not in gist.files:
                    continue

                if filename not in resources:
                    resources[filename] = ResourceFile.objects.create(project=project, file_name=filename, kind=kind,
                                                                      is_menu_icon=is_menu_icon)
                    # We already have this as a unicode string in .content, but it shouldn't have become unicode
                    # in the first place.
                    resources[filename].save_file(urllib2.urlopen(gist.files[filename].raw_url))
                ResourceIdentifier.objects.create(
                    resource_file=resources[filename],
                    resource_id=def_name,
                    character_regex=regex,
                    tracking=tracking
                )
        else:
            source_file = SourceFile.objects.create(project=project, file_name='app.js')
            source_file.save_file(gist.files['simply.js'].content)

    send_keen_event('cloudpebble', 'cloudpebble_gist_import', project=project, data={'data': {'gist_id': gist_id}})
    return project.id
<|endoftext|>


<commit_before>import imp
import os
import re
import time

from fabric.api import task, puts, env, local, sudo, run, execute, lcd
from fabric.tasks import Task
from fabric.utils import abort

try:
    import boto.ec2
    AWS_SUPPORT = True
    #import boto.ec2.elb
    #self.boto_ec2 = boto.ec2
    #self.boto_ec2_elb = boto.ec2.elb
except ImportError:
    AWS_SUPPORT = False

ip_re = re.compile('(?P<ip>\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})')

env.aptitude = "sudo aptitude -q -y"

def _get_ec2_connection():
    return boto.ec2.connect_to_region(
        os.environ.get('AWS_REGION', 'eu-west-1'))

def _find_ec2_by_name(name, state='running'):
    ec2 = _get_ec2_connection()
    reservations = ec2.get_all_instances(
        filters={'tag:Name': name, 'instance-state-name': state})
    if len(reservations) == 0:
        return None
    elif len(reservations) == 1:
        return reservations[0].instances[0]
    puts("Found more than one!")
    return None

class BaseHost(object):

    def __init__(self, settings):
        self.settings = settings

    def get_host_user(self):
        return getattr(self.settings, 'SSH_USER', env.user)

class EC2Host(BaseHost):

    def __init__(self, *args, **kwargs):
        super(EC2Host, self).__init__(*args, **kwargs)

    def get_ip(self):
        ec2 = _find_ec2_by_name(self.settings.EC2_NAME)
        return ec2.ip_address

class VagrantHost(BaseHost):

    def __init__(self, *args, **kwargs):
        super(VagrantHost, self).__init__(*args, **kwargs)
        self.vagrant_dir = os.path.dirname(self.settings.VAGRANT_FILE)

    def get_ip(self):
        ro = re.compile(r'^\s+config.+:private_network')
        with open(self.settings.VAGRANT_FILE, 'r') as vf:
            for line in vf:
                mo = ro.search(line)
                if mo is None:
                    continue
                mo = ip_re.search(line)
                if mo is not None:
                    return mo.group('ip')
        return None

class FabTaskIt(object):
    AWS_REQUIRED_VARS = [
        'AWS_ACCESS_KEY_ID',
        'AWS_AVAILABILITY_ZONE',
        'AWS_KEYPAIR_NAME',
        'AWS_REGION',
        'AWS_SECRET_ACCESS_KEY',
    ]

    def __init__(self):
        self.fabhosts = {}
        self.fabenvdirs = {}
        self.activated_hosts = []
        self.amis = {}

    def setup(self):
        fabhome = os.path.expanduser('~/.fab-task-it')

        if os.path.exists(fabhome):
            envdirshome = os.path.join(fabhome, 'envdirs')
            if os.path.exists(envdirshome):
                self.load_envdirs(envdirshome)

            hostshome = os.path.join(fabhome, 'hosts')
            if os.path.exists(hostshome):
                self.load_hosts(hostshome)

            ec2home = os.path.join(fabhome, 'ec2')
            if os.path.exists(ec2home):
                ami_list_file = os.path.join(ec2home, 'amis.py')
                if os.path.exists(ami_list_file):
                    self.load_amis(ami_list_file)

        self.setup_hosts()
        self.setup_environments()

    def load_envdirs(self, envdir):
        for envroot, envdirs, envfiles in os.walk(envdir):
            for d in envdirs:
                fabenvdir = {
                    'name': d,
                    'path': os.path.join(envdir, d),
                }
                fabenvdir['vars'] = self._get_env_vars(fabenvdir['path'])
                self.fabenvdirs[d] = fabenvdir

    def _get_env_vars(self, path):
        env_vars = {}
        for root, dirs, files in os.walk(path):
            for filename in files:
                with open(os.path.join(path, filename)) as f:
                    env_vars[filename] = f.readline().rstrip("\n")
        return env_vars

    def load_hosts(self, hostsdir):
        for root, dirs, files in os.walk(hostsdir):
            for filename in files:
                if filename[-3:] != '.py' or filename == '__init__.py':
                    continue
                modname = filename[:-3]
                fabhost = {'name': modname}
                try:
                    fabhost['settings'] = imp.load_source(
                        'fabhosts.{}'.format(modname),
                        os.path.join(hostsdir, filename)
                    )
                    self.fabhosts[fabhost['name']] = fabhost
                    if fabhost['settings'].HOST_TYPE == 'EC2':
                        fabhost['helper'] = EC2Host(fabhost['settings'])
                    elif fabhost['settings'].HOST_TYPE == 'VAGRANT':
                        fabhost['helper'] = VagrantHost(fabhost['settings'])
                except ImportError:
                    puts("host [{}] initialization failed.".format(filename))

    def load_amis(self, filename):
        try:
            self.amis_module = imp.load_source(
                'fabamis.amis',
                filename,
            )
            self.amis = self.amis_module.AMIS
        except ImportError:
            abort("amis file [{}] initialization failed.".format(filename))

    def load_environment(self, fabenv):
        envdir = self.fabenvdirs.get(fabenv, None)
        if envdir is None:
            puts('WARNING: {} environment not found!'.format(fabenv))
            return
        for k, v in envdir['vars'].items():
            os.environ[k] = v

    def load_host_environments(self, fabhost):
        for host_env in fabhost['settings'].ENVIRONMENTS:
            self.load_environment(host_env)
        env.user = self.get_host_user(fabhost['name'])

    def get_host_port(self, hostname=None):
        if hostname is None:
            hostname = self.activated_hosts[0]
        return getattr(self.fabhosts[hostname]['settings'], 'SSH_PORT', 22)

    def get_host_user(self, hostname=None):
        if hostname is None:
            hostname = self.activated_hosts[0]
        return self.fabhosts[hostname]['helper'].get_host_user()

    def activate_host(self, host):
        fabhost = self.fabhosts[host]
        self.load_host_environments(fabhost)
        env.hosts.append(fabhost['helper'].get_ip())
        self.activated_hosts.append(fabhost['name'])

    def get_activate_host_func(self, host):
        def f():
            return self.activate_host(host)
        return f

    def setup_hosts(self):
        for i, host in enumerate(self.fabhosts.keys()):
            globals()['host_{}'.format(i)] = SimpleTask(
                self.get_activate_host_func(host),
                name='host_{}'.format(host),
            )

    def activate_environment(self, fabenv):
        self.load_environment(fabenv)

    def get_activate_env_function(self, env):
        def f():
            return self.activate_environment(env)
        return f

    def setup_environments(self):
        for i, env in enumerate(self.fabenvdirs.keys()):
            globals()['env_{}'.format(i)] = SimpleTask(
                self.get_activate_env_function(env),
                name='env_{}'.format(self.fabenvdirs[env]['name']),
            )

    def get_available_files(self, *filenames):
        result = []
        for filename in filenames:
            pth = os.path.expanduser(filename)
            if os.path.exists(pth):
                result.append(pth)
        return result

    def get_ssh_private_key_files(self):
        return self.get_available_files('~/.ssh/id_dsa', '~/.ssh/id_rsa')

    def get_ssh_public_key_files(self):
        return self.get_available_files('~/.ssh/id_dsa.pub',
                                        '~/.ssh/id_rsa.pub')

    def get_ssh_command(self):
        port = self.get_host_port()
        user = self.get_host_user()
        host = env.hosts[0]
        hostname = self.activated_hosts[0]
        ssh_key_str = getattr(self.fabhosts[hostname]['settings'],
            'SSH_KEY', ' ')
        if ssh_key_str != ' ':
            ssh_key_str = ' -i {} '.format(ssh_key_str)
        return 'ssh -p {port} -l {user}{ssh_key}{host}'.format(
            host=host,
            port=port,
            user=user,
            ssh_key=ssh_key_str,
        )

    def get_active_host_helper(self, hostname=None):
        if hostname is None:
            hostname = self.activated_hosts[0]
        return self.fabhosts[hostname]['helper']

class SimpleTask(Task):
    def __init__(self, func, *args, **kwargs):
        super(SimpleTask, self).__init__(*args, **kwargs)
        self.func = func

    def run(self, *args, **kwargs):
        return self.func(*args, **kwargs)

class AWSTask(Task):
    def __init__(self, func, *args, **kwargs):
        super(AWSTask, self).__init__(*args, **kwargs)
        self.func = func

    def run(self, *args, **kwargs):
        if not AWS_SUPPORT:
            abort('No aws support! Install boto.')
        for envvar in fabtaskit.AWS_REQUIRED_VARS:
            if envvar not in os.environ:
                abort('AWS environment not configured properly!')
        return self.func(*args, **kwargs)

fabtaskit = FabTaskIt()
fabtaskit.setup()

def print_ec2(ec2):
    title = "Name: {}".format(ec2.tags['Name'])
    puts(title)
    puts("=" * len(title))
    puts("ID: {}".format(ec2.id))
    puts("Name: {}".format(ec2.tags['Name']))
    puts("State: {}".format(ec2.state))
    puts("Type: {}".format(ec2.instance_type))
    puts("Public DNS: {}".format(ec2.dns_name))
    puts("Public IP: {}".format(ec2.ip_address))
    puts("Private DNS: {}".format(ec2.private_dns_name))
    puts("Private IP: {}".format(ec2.private_ip_address))
    # http://docs.pythonboto.org/en/latest/ref/ec2.html \
    #   #boto.ec2.instance.Instance.get_attribute
    ia = ec2.get_attribute('disableApiTermination')
    puts("Termination Protection: {}".format(ia['disableApiTermination']))
    ia = ec2.get_attribute('instanceInitiatedShutdownBehavior')
    puts("Shutdown Behavior: {}".format(
        ia['instanceInitiatedShutdownBehavior']))
    puts("")

@task
def test_hostname():
    puts(fabtaskit.activated_hosts)

@task
def test_environment():
    puts(os.environ)

@task(task_class=AWSTask)
def find_ec2_by_name(name):
    ec2 = _find_ec2_by_name(name)
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    print_ec2(ec2)

@task
def login():
    local(fabtaskit.get_ssh_command())

@task
def ec2_list_all():
    ec2 = _get_ec2_connection()
    instances = ec2.get_all_instances()
    for reservation in instances:
        print_ec2(reservation.instances[0])

@task
def supervisor_restart(service):
    sudo('supervisorctl restart {0}'.format(service))

@task
def supervisor_status():
    sudo('supervisorctl status')

@task
def tailf(filepath):
    env.output_prefix = False
    sudo('tail -f {0}'.format(filepath))
    env.output_prefix = True

@task
def copy_ssh_pub_keys():
    keys = []

    for filename in fabtaskit.get_ssh_public_key_files():
        try:
            with open(filename, 'r') as f:
                keys.append(f.readline().rstrip("\n"))
        except IOError:
            pass

    puts('Got {} keys'.format(len(keys)))

    ssh_cmd = fabtaskit.get_ssh_command()

    result = local(
        "{ssh_cmd} 'cat .ssh/authorized_keys'".format(ssh_cmd=ssh_cmd),
        capture=True,
    )
    for key in keys:
        if result.find(key) < 0:
            puts('not found key:')
            puts(key)
            local("""{ssh_cmd} 'echo "{key}" >> .ssh/authorized_keys'"""
                  .format(ssh_cmd=ssh_cmd, key=key))

@task
def aptitude_update():
    run("{0} update".format(env.aptitude), quiet=True)

@task
def aptitude_safe_upgrade():
    execute(aptitude_update)
    run("{0} safe-upgrade".format(env.aptitude))

@task
def vagrant_up():
    vagrant = fabtaskit.get_active_host_helper()
    with lcd(vagrant.vagrant_dir):
        local('vagrant up')

@task
def vagrant_halt():
    vagrant = fabtaskit.get_active_host_helper()
    with lcd(vagrant.vagrant_dir):
        local('vagrant halt')

@task(task_class=AWSTask)
def ec2_terminate(name):
    ec2 = _find_ec2_by_name(name)
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    try:
        ec2.terminate()
        puts("Instance terminated!!")
    except boto.exception.EC2ResponseError as ex:
        puts("[{}] {}".format(ex.error_code, ex.message))

@task(task_class=AWSTask)
def ec2_run_instances(ami_name, tag_name):
    ec2 = _get_ec2_connection()
    ami = fabtaskit.amis.get(ami_name)
    if ami is None:
        abort("Couldn't find ami conf with name: {}".format(ami_name))
    reservation = ec2.run_instances(
        image_id=ami['image_id'],
        key_name=ami['key_name'],
        security_groups=ami['security_groups'],
        instance_type=ami['instance_type'],
        placement=ami['placement'],
        instance_initiated_shutdown_behavior=ami[
            'instance_initiated_shutdown_behavior'],
    )
    instance = reservation.instances[0]
    ec2.create_tags([instance.id], {'Name': tag_name})
    tries = 5
    booted = False
    while not booted and tries != 0:
        time.sleep(5)
        for r in ec2.get_all_instances():
            if r.id == reservation.id:
                print_ec2(r.instances[0])
                booted = True
                break
        tries -= 1
    if not booted:
        abort("Couldn't find if instance booted")

@task(task_class=AWSTask)
def ec2_stop_instance(name):
    ec2 = _find_ec2_by_name(name)
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    try:
        ec2.stop()
        puts("Instance stopped!!")
    except boto.exception.EC2ResponseError as ex:
        puts("[{}] {}".format(ex.error_code, ex.message))

@task(task_class=AWSTask)
def ec2_start_instance(name):
    ec2 = _find_ec2_by_name(name, state='stopped')
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    try:
        ec2.start()
        puts("Instance started!!")
    except boto.exception.EC2ResponseError as ex:
        puts("[{}] {}".format(ex.error_code, ex.message))
<commit_msg>changed method name to ec2_find_by_name. Using AWS keys from environment<commit_after>import imp
import os
import re
import time

from fabric.api import task, puts, env, local, sudo, run, execute, lcd
from fabric.tasks import Task
from fabric.utils import abort

try:
    import boto.ec2
    AWS_SUPPORT = True
    #import boto.ec2.elb
    #self.boto_ec2 = boto.ec2
    #self.boto_ec2_elb = boto.ec2.elb
except ImportError:
    AWS_SUPPORT = False

ip_re = re.compile('(?P<ip>\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})')

env.aptitude = "sudo aptitude -q -y"

def _get_ec2_connection():
    return boto.ec2.connect_to_region(
        os.environ.get('AWS_REGION', 'eu-west-1'))

def _find_ec2_by_name(name, state='running'):
    ec2 = _get_ec2_connection()
    reservations = ec2.get_all_instances(
        filters={'tag:Name': name, 'instance-state-name': state})
    if len(reservations) == 0:
        return None
    elif len(reservations) == 1:
        return reservations[0].instances[0]
    puts("Found more than one!")
    return None

class BaseHost(object):

    def __init__(self, settings):
        self.settings = settings

    def get_host_user(self):
        return getattr(self.settings, 'SSH_USER', env.user)

class EC2Host(BaseHost):

    def __init__(self, *args, **kwargs):
        super(EC2Host, self).__init__(*args, **kwargs)

    def get_ip(self):
        ec2 = _find_ec2_by_name(self.settings.EC2_NAME)
        return ec2.ip_address

class VagrantHost(BaseHost):

    def __init__(self, *args, **kwargs):
        super(VagrantHost, self).__init__(*args, **kwargs)
        self.vagrant_dir = os.path.dirname(self.settings.VAGRANT_FILE)

    def get_ip(self):
        ro = re.compile(r'^\s+config.+:private_network')
        with open(self.settings.VAGRANT_FILE, 'r') as vf:
            for line in vf:
                mo = ro.search(line)
                if mo is None:
                    continue
                mo = ip_re.search(line)
                if mo is not None:
                    return mo.group('ip')
        return None

class FabTaskIt(object):
    AWS_REQUIRED_VARS = [
        'AWS_ACCESS_KEY_ID',
        'AWS_AVAILABILITY_ZONE',
        'AWS_KEYPAIR_NAME',
        'AWS_REGION',
        'AWS_SECRET_ACCESS_KEY',
    ]

    def __init__(self):
        self.fabhosts = {}
        self.fabenvdirs = {}
        self.activated_hosts = []
        self.amis = {}

    def setup(self):
        fabhome = os.path.expanduser('~/.fab-task-it')

        if os.path.exists(fabhome):
            envdirshome = os.path.join(fabhome, 'envdirs')
            if os.path.exists(envdirshome):
                self.load_envdirs(envdirshome)

            hostshome = os.path.join(fabhome, 'hosts')
            if os.path.exists(hostshome):
                self.load_hosts(hostshome)

            ec2home = os.path.join(fabhome, 'ec2')
            if os.path.exists(ec2home):
                ami_list_file = os.path.join(ec2home, 'amis.py')
                if os.path.exists(ami_list_file):
                    self.load_amis(ami_list_file)

        self.setup_hosts()
        self.setup_environments()

    def load_envdirs(self, envdir):
        for envroot, envdirs, envfiles in os.walk(envdir):
            for d in envdirs:
                fabenvdir = {
                    'name': d,
                    'path': os.path.join(envdir, d),
                }
                fabenvdir['vars'] = self._get_env_vars(fabenvdir['path'])
                self.fabenvdirs[d] = fabenvdir

    def _get_env_vars(self, path):
        env_vars = {}
        for root, dirs, files in os.walk(path):
            for filename in files:
                with open(os.path.join(path, filename)) as f:
                    env_vars[filename] = f.readline().rstrip("\n")
        return env_vars

    def load_hosts(self, hostsdir):
        for root, dirs, files in os.walk(hostsdir):
            for filename in files:
                if filename[-3:] != '.py' or filename == '__init__.py':
                    continue
                modname = filename[:-3]
                fabhost = {'name': modname}
                try:
                    fabhost['settings'] = imp.load_source(
                        'fabhosts.{}'.format(modname),
                        os.path.join(hostsdir, filename)
                    )
                    self.fabhosts[fabhost['name']] = fabhost
                    if fabhost['settings'].HOST_TYPE == 'EC2':
                        fabhost['helper'] = EC2Host(fabhost['settings'])
                    elif fabhost['settings'].HOST_TYPE == 'VAGRANT':
                        fabhost['helper'] = VagrantHost(fabhost['settings'])
                except ImportError:
                    puts("host [{}] initialization failed.".format(filename))

    def load_amis(self, filename):
        try:
            self.amis_module = imp.load_source(
                'fabamis.amis',
                filename,
            )
            self.amis = self.amis_module.AMIS
        except ImportError:
            abort("amis file [{}] initialization failed.".format(filename))

    def load_environment(self, fabenv):
        envdir = self.fabenvdirs.get(fabenv, None)
        if envdir is None:
            puts('WARNING: {} environment not found!'.format(fabenv))
            return
        for k, v in envdir['vars'].items():
            os.environ[k] = v

    def load_host_environments(self, fabhost):
        for host_env in fabhost['settings'].ENVIRONMENTS:
            self.load_environment(host_env)
        env.user = self.get_host_user(fabhost['name'])

    def get_host_port(self, hostname=None):
        if hostname is None:
            hostname = self.activated_hosts[0]
        return getattr(self.fabhosts[hostname]['settings'], 'SSH_PORT', 22)

    def get_host_user(self, hostname=None):
        if hostname is None:
            hostname = self.activated_hosts[0]
        return self.fabhosts[hostname]['helper'].get_host_user()

    def activate_host(self, host):
        fabhost = self.fabhosts[host]
        self.load_host_environments(fabhost)
        env.hosts.append(fabhost['helper'].get_ip())
        self.activated_hosts.append(fabhost['name'])

    def get_activate_host_func(self, host):
        def f():
            return self.activate_host(host)
        return f

    def setup_hosts(self):
        for i, host in enumerate(self.fabhosts.keys()):
            globals()['host_{}'.format(i)] = SimpleTask(
                self.get_activate_host_func(host),
                name='host_{}'.format(host),
            )

    def activate_environment(self, fabenv):
        self.load_environment(fabenv)

    def get_activate_env_function(self, env):
        def f():
            return self.activate_environment(env)
        return f

    def setup_environments(self):
        for i, env in enumerate(self.fabenvdirs.keys()):
            globals()['env_{}'.format(i)] = SimpleTask(
                self.get_activate_env_function(env),
                name='env_{}'.format(self.fabenvdirs[env]['name']),
            )

    def get_available_files(self, *filenames):
        result = []
        for filename in filenames:
            pth = os.path.expanduser(filename)
            if os.path.exists(pth):
                result.append(pth)
        return result

    def get_ssh_private_key_files(self):
        return self.get_available_files('~/.ssh/id_dsa', '~/.ssh/id_rsa')

    def get_ssh_public_key_files(self):
        return self.get_available_files('~/.ssh/id_dsa.pub',
                                        '~/.ssh/id_rsa.pub')

    def get_ssh_command(self):
        port = self.get_host_port()
        user = self.get_host_user()
        host = env.hosts[0]
        hostname = self.activated_hosts[0]
        ssh_key_str = getattr(self.fabhosts[hostname]['settings'],
            'SSH_KEY', ' ')
        if ssh_key_str != ' ':
            ssh_key_str = ' -i {} '.format(ssh_key_str)
        return 'ssh -p {port} -l {user}{ssh_key}{host}'.format(
            host=host,
            port=port,
            user=user,
            ssh_key=ssh_key_str,
        )

    def get_active_host_helper(self, hostname=None):
        if hostname is None:
            hostname = self.activated_hosts[0]
        return self.fabhosts[hostname]['helper']

class SimpleTask(Task):
    def __init__(self, func, *args, **kwargs):
        super(SimpleTask, self).__init__(*args, **kwargs)
        self.func = func

    def run(self, *args, **kwargs):
        return self.func(*args, **kwargs)

class AWSTask(Task):
    def __init__(self, func, *args, **kwargs):
        super(AWSTask, self).__init__(*args, **kwargs)
        self.func = func

    def run(self, *args, **kwargs):
        if not AWS_SUPPORT:
            abort('No aws support! Install boto.')
        for envvar in fabtaskit.AWS_REQUIRED_VARS:
            if envvar not in os.environ:
                abort('AWS environment not configured properly!')
        return self.func(*args, **kwargs)

fabtaskit = FabTaskIt()
fabtaskit.setup()

def print_ec2(ec2):
    title = "Name: {}".format(ec2.tags['Name'])
    puts(title)
    puts("=" * len(title))
    puts("ID: {}".format(ec2.id))
    puts("Name: {}".format(ec2.tags['Name']))
    puts("State: {}".format(ec2.state))
    puts("Type: {}".format(ec2.instance_type))
    puts("Public DNS: {}".format(ec2.dns_name))
    puts("Public IP: {}".format(ec2.ip_address))
    puts("Private DNS: {}".format(ec2.private_dns_name))
    puts("Private IP: {}".format(ec2.private_ip_address))
    # http://docs.pythonboto.org/en/latest/ref/ec2.html \
    #   #boto.ec2.instance.Instance.get_attribute
    ia = ec2.get_attribute('disableApiTermination')
    puts("Termination Protection: {}".format(ia['disableApiTermination']))
    ia = ec2.get_attribute('instanceInitiatedShutdownBehavior')
    puts("Shutdown Behavior: {}".format(
        ia['instanceInitiatedShutdownBehavior']))
    puts("")

@task
def test_hostname():
    puts(fabtaskit.activated_hosts)

@task
def test_environment():
    puts(os.environ)

@task(task_class=AWSTask)
def ec2_find_by_name(name):
    ec2 = _find_ec2_by_name(name)
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    print_ec2(ec2)

@task
def login():
    local(fabtaskit.get_ssh_command())

@task
def ec2_list_all():
    ec2 = _get_ec2_connection()
    instances = ec2.get_all_instances()
    for reservation in instances:
        print_ec2(reservation.instances[0])

@task
def supervisor_restart(service):
    sudo('supervisorctl restart {0}'.format(service))

@task
def supervisor_status():
    sudo('supervisorctl status')

@task
def tailf(filepath):
    env.output_prefix = False
    sudo('tail -f {0}'.format(filepath))
    env.output_prefix = True

@task
def copy_ssh_pub_keys():
    keys = []

    for filename in fabtaskit.get_ssh_public_key_files():
        try:
            with open(filename, 'r') as f:
                keys.append(f.readline().rstrip("\n"))
        except IOError:
            pass

    puts('Got {} keys'.format(len(keys)))

    ssh_cmd = fabtaskit.get_ssh_command()

    result = local(
        "{ssh_cmd} 'cat .ssh/authorized_keys'".format(ssh_cmd=ssh_cmd),
        capture=True,
    )
    for key in keys:
        if result.find(key) < 0:
            puts('not found key:')
            puts(key)
            local("""{ssh_cmd} 'echo "{key}" >> .ssh/authorized_keys'"""
                  .format(ssh_cmd=ssh_cmd, key=key))

@task
def aptitude_update():
    run("{0} update".format(env.aptitude), quiet=True)

@task
def aptitude_safe_upgrade():
    execute(aptitude_update)
    run("{0} safe-upgrade".format(env.aptitude))

@task
def vagrant_up():
    vagrant = fabtaskit.get_active_host_helper()
    with lcd(vagrant.vagrant_dir):
        local('vagrant up')

@task
def vagrant_halt():
    vagrant = fabtaskit.get_active_host_helper()
    with lcd(vagrant.vagrant_dir):
        local('vagrant halt')

@task(task_class=AWSTask)
def ec2_terminate(name):
    ec2 = _find_ec2_by_name(name)
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    try:
        ec2.terminate()
        puts("Instance terminated!!")
    except boto.exception.EC2ResponseError as ex:
        puts("[{}] {}".format(ex.error_code, ex.message))

@task(task_class=AWSTask)
def ec2_run_instances(ami_name, tag_name):
    ec2 = _get_ec2_connection()
    ami = fabtaskit.amis.get(ami_name)
    if ami is None:
        abort("Couldn't find ami conf with name: {}".format(ami_name))
    kwargs = ami.copy()
    kwargs['key_name'] = os.environ['AWS_KEYPAIR_NAME']
    reservation = ec2.run_instances(**kwargs)
    instance = reservation.instances[0]
    ec2.create_tags([instance.id], {'Name': tag_name})
    tries = 5
    booted = False
    while not booted and tries != 0:
        time.sleep(5)
        for r in ec2.get_all_instances():
            if r.id == reservation.id:
                print_ec2(r.instances[0])
                booted = True
                break
        tries -= 1
    if not booted:
        abort("Couldn't find if instance booted")

@task(task_class=AWSTask)
def ec2_stop_instance(name):
    ec2 = _find_ec2_by_name(name)
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    try:
        ec2.stop()
        puts("Instance stopped!!")
    except boto.exception.EC2ResponseError as ex:
        puts("[{}] {}".format(ex.error_code, ex.message))

@task(task_class=AWSTask)
def ec2_start_instance(name):
    ec2 = _find_ec2_by_name(name, state='stopped')
    if ec2 is None:
        abort("Couldn't find '{}'".format(name))
    try:
        ec2.start()
        puts("Instance started!!")
    except boto.exception.EC2ResponseError as ex:
        puts("[{}] {}".format(ex.error_code, ex.message))
<|endoftext|>