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

"""
This executable module is a console application for presenting itself as a
MODBUS server accepting read and write MODBUS PDUs.
"""

import os
import logging

from bacpypes.debugging import bacpypes_debugging, ModuleLogger
from bacpypes.consolecmd import ConsoleCmd
from bacpypes.consolelogging import ArgumentParser

from bacpypes.comm import Client, bind
from bacpypes.core import run

from .pdu import ExceptionResponse, \
    ReadCoilsResponse, ReadDiscreteInputsResponse, ReadMultipleRegistersResponse, \
    WriteSingleCoilResponse, WriteSingleRegisterResponse, WriteMultipleRegistersResponse
from .app import ModbusServer, ModbusException

# some debugging
_debug = 0
_log = ModuleLogger(globals())
_commlog = logging.getLogger(__name__ + "._commlog")

# settings
SERVER_HOST = os.getenv("SERVER_HOST", "")
SERVER_PORT = int(os.getenv("SERVER_PORT", 502))

#
#   ConsoleServer
#

@bacpypes_debugging
class ConsoleServer(ConsoleCmd, Client):

    """
    Console Server
    """

    def __init__(self):
        if _debug: ConsoleServer._debug("__init__")
        ConsoleCmd.__init__(self)

        # create some coils and registers
        self.coils = [False] * 10
        self.registers = [0] * 10

    def confirmation(self, req):
        """Got a request from a client."""
        if _debug: ConsoleServer._debug("confirmation %r", req)
        _commlog.debug(">>> %r %r", req.pduSource, req)

        try:
            # look up a matching function
            try:
                fn = getattr(self, "do_" + req.__class__.__name__)
            except AttributeError:
                raise ModbusException(ExceptionResponse.ILLEGAL_FUNCTION)

            # try to execute it
            resp = fn(req)

        except ModbusException as err:
            # create an exception response
            resp = ExceptionResponse(req.mpduFunctionCode, err.errCode)

        # match the transaction information
        resp.pduDestination = req.pduSource
        resp.mpduTransactionID = req.mpduTransactionID
        _commlog.debug("<<< %r %r", resp.pduDestination, resp)

        # send the response back
        self.request(resp)

    def pull_coils(self, address, count):
        """Called when there is a request for the current value of a coil."""
        if _debug: ConsoleServer._debug("pull_coils %r %r", address, count)

    def push_coils(self, address, count):
        """Called when a MODBUS service has changed the value of one or more coils."""
        if _debug: ConsoleServer._debug("push_coils %r %r", address, count)

    def pull_registers(self, address, count):
        """Called when a MODBUS client is requesting the current value of one
        or more registers."""
        if _debug: ConsoleServer._debug("pull_registers %r %r", address, count)

    def push_registers(self, address, count):
        """Called when a MODBUS service has changed the value of one or more
        registers."""
        if _debug: ConsoleServer._debug("push_registers %r %r", address, count)

    # ---------- Coils ----------

    def do_ReadCoilsRequest(self, req):
        ConsoleServer._debug('do_ReadCoilsRequest %r', req)
        if (req.address + req.count) > len(self.coils):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_coils(req.address, req.count)

        return ReadCoilsResponse(self.coils[req.address:req.address+req.count])

    def do_WriteSingleCoilRequest(self, req):
        ConsoleServer._debug('do_WriteSingleCoilRequest %r', req)
        if req.address > len(self.coils):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        # check the value and save it
        if (req.value == 0x00):
            self.coils[req.address] = 0
        elif (req.value == 0xFF):
            self.coils[req.address] = 1
        else:
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_VALUE)

        self.push_coils(req.address, 1)

        # return the new value
        return WriteSingleCoilResponse(req.address, req.value)

    # ---------- Descrete Inputs (mapped as a coil) ----------

    def do_ReadDescreteInputsRequest(self, req):
        ConsoleServer._debug('do_ReadDescreteInputsRequest %r', req)
        if (req.address + req.count) > len(self.coils):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_coils(req.address, req.count)

        return ReadDiscreteInputsResponse(self.coils[req.address:req.address+req.count])

    # ---------- Registers ----------

    def do_ReadMultipleRegistersRequest(self, req):
        ConsoleServer._debug('do_ReadMultipleRegistersRequest %r', req)
        if (req.address + req.count) > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_registers(req.address, req.count)

        return ReadMultipleRegistersResponse(self.registers[req.address:req.address+req.count])

    def do_WriteSingleRegisterRequest(self, req):
        ConsoleServer._debug('do_WriteSingleRegisterRequest %r', req)
        if req.address > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        # save the value
        self.registers[req.address] = req.value

        self.push_registers(req.address, 1)

        # return the new value
        return WriteSingleRegisterResponse(req.address, req.value)

    def do_WriteMultipleRegistersRequest(self, req):
        ConsoleServer._debug('do_WriteMultipleRegistersRequest %r', req)
        if (req.address + req.count) > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        # save the values
        for i in range(req.count):
            self.registers[req.address + i] = req.registers[i]

        self.push_registers(req.address, req.count)

        return WriteMultipleRegistersResponse(req.address, req.count)

    # ---------- Input Registers (mapped as a register) ----------

    def do_ReadInputRegistersRequest(self, req):
        ConsoleServer._debug('do_ReadInputRegistersRequest %r', req)
        if (req.address + req.count) > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_registers(req.address, req.count)

        return ReadMultipleRegistersResponse(self.registers[req.address:req.address+req.count])

#
#   main
#

def main():
    # parse the command line arguments
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        "--host", type=str,
        help="address of host (default {!r})".format(SERVER_HOST),
        default=SERVER_HOST,
        )
    parser.add_argument(
        "--port", type=int,
        help="server port (default {!r})".format(SERVER_PORT),
        default=SERVER_PORT,
        )
    args = parser.parse_args()

    if _debug: _log.debug("initialization")
    if _debug: _log.debug("    - args: %r", args)

    # local IO functions
    bind(ConsoleServer(), ModbusServer(port=args.port))

    _log.debug("running")

    run()

    _log.debug("fini")

if __name__ == "__main__":
    main()

<commit_msg>rename the server, pull out the console part (maybe it goes back in later)<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This executable module is a console application for presenting itself as a
MODBUS server accepting read and write MODBUS PDUs.
"""

import os
import logging

from bacpypes.debugging import bacpypes_debugging, ModuleLogger
from bacpypes.consolecmd import ConsoleCmd
from bacpypes.consolelogging import ArgumentParser

from bacpypes.comm import Client, bind
from bacpypes.core import run

from .pdu import ExceptionResponse, \
    ReadCoilsResponse, ReadDiscreteInputsResponse, ReadMultipleRegistersResponse, \
    WriteSingleCoilResponse, WriteSingleRegisterResponse, WriteMultipleRegistersResponse
from .app import ModbusServer, ModbusException

# some debugging
_debug = 0
_log = ModuleLogger(globals())
_commlog = logging.getLogger(__name__ + "._commlog")

# settings
SERVER_HOST = os.getenv("SERVER_HOST", "")
SERVER_PORT = int(os.getenv("SERVER_PORT", 502))

#
#   SimpleServer
#

@bacpypes_debugging
class SimpleServer(Client):

    """
    Simple Server
    """

    def __init__(self, unitNumber=1):
        if _debug: SimpleServer._debug("__init__")
        Client.__init__(self)

        # save the unit number
        self.unitNumber = unitNumber

        # create some coils and registers
        self.coils = [False] * 10
        self.registers = [0] * 10

    def confirmation(self, req):
        """Got a request from a client."""
        if _debug: SimpleServer._debug("confirmation %r", req)
        _commlog.debug(">>> %r %r", req.pduSource, req)

        # if its an exception, punt
        if isinstance(req, Exception):
            if _debug: SimpleServer._debug("    - punt exceptions")
            return

        # if it's not for us, dump it
        if req.mpduUnitID != self.unitNumber:
            if _debug: SimpleServer._debug("    - not for us")
            return

        try:
            # look up a matching function
            try:
                fn = getattr(self, "do_" + req.__class__.__name__)
            except AttributeError:
                raise ModbusException(ExceptionResponse.ILLEGAL_FUNCTION)

            # try to execute it
            resp = fn(req)

        except ModbusException as err:
            # create an exception response
            resp = ExceptionResponse(req.mpduFunctionCode, err.errCode)

        # match the transaction information
        resp.pduDestination = req.pduSource
        resp.mpduTransactionID = req.mpduTransactionID
        _commlog.debug("<<< %r %r", resp.pduDestination, resp)

        # send the response back
        self.request(resp)

    def pull_coils(self, address, count):
        """Called when there is a request for the current value of a coil."""
        if _debug: SimpleServer._debug("pull_coils %r %r", address, count)

    def push_coils(self, address, count):
        """Called when a MODBUS service has changed the value of one or more coils."""
        if _debug: SimpleServer._debug("push_coils %r %r", address, count)

    def pull_registers(self, address, count):
        """Called when a MODBUS client is requesting the current value of one
        or more registers."""
        if _debug: SimpleServer._debug("pull_registers %r %r", address, count)

    def push_registers(self, address, count):
        """Called when a MODBUS service has changed the value of one or more
        registers."""
        if _debug: SimpleServer._debug("push_registers %r %r", address, count)

    # ---------- Coils ----------

    def do_ReadCoilsRequest(self, req):
        SimpleServer._debug('do_ReadCoilsRequest %r', req)
        if (req.address + req.count) > len(self.coils):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_coils(req.address, req.count)

        return ReadCoilsResponse(self.coils[req.address:req.address+req.count])

    def do_WriteSingleCoilRequest(self, req):
        SimpleServer._debug('do_WriteSingleCoilRequest %r', req)
        if req.address > len(self.coils):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        # check the value and save it
        if (req.value == 0x00):
            self.coils[req.address] = 0
        elif (req.value == 0xFF):
            self.coils[req.address] = 1
        else:
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_VALUE)

        self.push_coils(req.address, 1)

        # return the new value
        return WriteSingleCoilResponse(req.address, req.value)

    # ---------- Descrete Inputs (mapped as a coil) ----------

    def do_ReadDescreteInputsRequest(self, req):
        SimpleServer._debug('do_ReadDescreteInputsRequest %r', req)
        if (req.address + req.count) > len(self.coils):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_coils(req.address, req.count)

        return ReadDiscreteInputsResponse(self.coils[req.address:req.address+req.count])

    # ---------- Registers ----------

    def do_ReadMultipleRegistersRequest(self, req):
        SimpleServer._debug('do_ReadMultipleRegistersRequest %r', req)
        if (req.address + req.count) > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_registers(req.address, req.count)

        return ReadMultipleRegistersResponse(self.registers[req.address:req.address+req.count])

    def do_WriteSingleRegisterRequest(self, req):
        SimpleServer._debug('do_WriteSingleRegisterRequest %r', req)
        if req.address > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        # save the value
        self.registers[req.address] = req.value

        self.push_registers(req.address, 1)

        # return the new value
        return WriteSingleRegisterResponse(req.address, req.value)

    def do_WriteMultipleRegistersRequest(self, req):
        SimpleServer._debug('do_WriteMultipleRegistersRequest %r', req)
        if (req.address + req.count) > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        # save the values
        for i in range(req.count):
            self.registers[req.address + i] = req.registers[i]

        self.push_registers(req.address, req.count)

        return WriteMultipleRegistersResponse(req.address, req.count)

    # ---------- Input Registers (mapped as a register) ----------

    def do_ReadInputRegistersRequest(self, req):
        SimpleServer._debug('do_ReadInputRegistersRequest %r', req)
        if (req.address + req.count) > len(self.registers):
            raise ModbusException(ExceptionResponse.ILLEGAL_DATA_ADDRESS)

        self.pull_registers(req.address, req.count)

        return ReadMultipleRegistersResponse(self.registers[req.address:req.address+req.count])

#
#   main
#

def main():
    # parse the command line arguments
    parser = ArgumentParser(description=__doc__)
    parser.add_argument(
        "--host", type=str,
        help="address of host (default {!r})".format(SERVER_HOST),
        default=SERVER_HOST,
        )
    parser.add_argument(
        "--port", type=int,
        help="server port (default {!r})".format(SERVER_PORT),
        default=SERVER_PORT,
        )
    args = parser.parse_args()

    if _debug: _log.debug("initialization")
    if _debug: _log.debug("    - args: %r", args)

    # local IO functions
    bind(SimpleServer(), ModbusServer(port=args.port))

    _log.debug("running")

    run()

    _log.debug("fini")

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>import requests
import regex
import subprocess
import os
import shutil

__author__ = 'gpchelkin'

token = os.getenv('TG_BOT_TOKEN')
apiurl = 'https://api.telegram.org/bot' + token + '/'
offset = 0

homedir = os.path.expanduser('~')

scdlconfdir = os.path.join(homedir,'.config/scdl')
if os.path.exists(scdlconfdir):
    os.rmdir(scdlconfdir)
os.makedirs(scdlconfdir)
shutil.copy('scdl.cfg',scdlconfdir)

scdldir = os.path.join(homedir,'scdldir')
if os.path.exists(scdldir):
    os.rmdir(scdldir)

scdlbin = 'scdl -l '
scdlopts = ' -c --path ' + scdldir + ' --onlymp3 --addtofile --hide-progress --hidewarnings'

def sendaudio(scdlurl, chat_id):
    rmsg = requests.post(apiurl + 'sendMessage', json=dict(chat_id=chat_id, parse_mode='Markdown', text='_Wait a bit, downloading and sending..._'))
    print(scdldir)
    os.makedirs(scdldir)
    print("1")
    subprocess.call(scdlbin + scdlurl + scdlopts, shell=True)
    print("2")
    scdlfile = os.listdir(scdldir)[0]
    print(os.listdir(scdldir))
    scdlfullpath = os.path.join(scdldir,scdlfile)
    print(scdlfullpath)
    trackopen = open(scdlfullpath, 'rb')
    print(trackopen)
    track = trackopen.read()
    print(track)
    raudio = requests.post(apiurl + 'sendAudio',
                           files=dict(audio=(scdlfile, track, 'audio/mpeg')),
                           data=dict(chat_id=chat_id))
    print('audio ' + scdlfile + ' sent')
    print(raudio.json()['result'])
    os.rmdir(scdldir)

while 1:
    r = requests.post(apiurl + 'getUpdates', json=dict(offset=offset, timeout=60))
    try:
        updates = r.json()['result']
        for i in range(len(updates)):
            offset = updates[i]['update_id']
            message = updates[i].get('message')
            if message:
                chat_id = message['chat']['id']
                gottext = message.get('text')
                if gottext:
                    print('got text ' + gottext)  # TODO log time and all to file
                    if regex.compile('^/start(@scdlbot)?(\s+.*)?$').match(gottext.strip()) or \
                            regex.compile('^/help(@scdlbot)?(\s+.*)?$').match(gottext.strip()):
                        text = open('help.md', 'r').read()
                        rmsg = requests.post(apiurl + 'sendMessage', json=dict(chat_id=chat_id, text=text, parse_mode='Markdown',
                                                                               disable_web_page_preview=True))
                    elif regex.compile('^/dl(@scdlbot)?(\s+.*)?$').match(gottext.strip()):
                        sendaudio(gottext.split()[1],chat_id)
                    elif regex.compile('^http(s)?://(m.)?soundcloud.com/([\w\d_-]+)/([\w\d_-]+)$').match(gottext.strip()):
                        sendaudio(gottext,chat_id)
                else:
                    print('got idkwhat')
                msgfrom = message.get('from')
                print('  from ' + msgfrom['first_name'] + ' ' + msgfrom.get('last_name', '') +
                          ' @' + msgfrom.get('username', ''))

    except (KeyboardInterrupt, SystemExit):
        raise

    except:
        print('smth went wrong')
        #no key 'result' in dictionary -> very bad
        #if r.status_code == requests.codes.ok:
        #if r.json():
        #if r.json()['ok']:
        #if len(updates) != 0:
        #if msgfrom

    finally:
        offset += 1
<commit_msg>More debug print<commit_after>import requests
import regex
import subprocess
import os
import shutil

__author__ = 'gpchelkin'

token = os.getenv('TG_BOT_TOKEN')
apiurl = 'https://api.telegram.org/bot' + token + '/'
offset = 0

homedir = os.path.expanduser('~')

scdlconfdir = os.path.join(homedir,'.config/scdl')
if os.path.exists(scdlconfdir):
    shutil.rmtree(scdlconfdir)
os.makedirs(scdlconfdir)
shutil.copy('scdl.cfg',scdlconfdir)

scdldir = os.path.join(homedir,'scdldir')
if os.path.exists(scdldir):
    shutil.rmtree(scdldir)

scdlbin = 'scdl -l '
scdlopts = ' -c --path ' + scdldir + ' --onlymp3 --addtofile --hide-progress --hidewarnings'

def sendaudio(scdlurl, chat_id):
    rmsg = requests.post(apiurl + 'sendMessage', json=dict(chat_id=chat_id, parse_mode='Markdown', text='_Wait a bit, downloading and sending..._'))
    if os.path.exists(scdldir):
        shutil.rmtree(scdldir)
    print("0")
    os.makedirs(scdldir)
    print("1")
    subprocess.call(scdlbin + scdlurl + scdlopts, shell=True)
    print("2")
    scdlfile = os.listdir(scdldir)[0]
    print(os.listdir(scdldir))
    scdlfullpath = os.path.join(scdldir,scdlfile)
    print(scdlfullpath)
    trackopen = open(scdlfullpath, 'rb')
    print(trackopen)
    track = trackopen.read()
    print(track)
    raudio = requests.post(apiurl + 'sendAudio',
                           files=dict(audio=(scdlfile, track, 'audio/mpeg')),
                           data=dict(chat_id=chat_id))
    print('audio ' + scdlfile + ' sent')
    print(raudio.json()['result'])
    shutil.rmtree(scdldir)

while 1:
    r = requests.post(apiurl + 'getUpdates', json=dict(offset=offset, timeout=60))
    try:
        updates = r.json()['result']
        for i in range(len(updates)):
            offset = updates[i]['update_id']
            message = updates[i].get('message')
            if message:
                chat_id = message['chat']['id']
                gottext = message.get('text')
                if gottext:
                    print('got text ' + gottext)  # TODO log time and all to file
                    if regex.compile('^/start(@scdlbot)?(\s+.*)?$').match(gottext.strip()) or \
                            regex.compile('^/help(@scdlbot)?(\s+.*)?$').match(gottext.strip()):
                        text = open('help.md', 'r').read()
                        rmsg = requests.post(apiurl + 'sendMessage', json=dict(chat_id=chat_id, text=text, parse_mode='Markdown',
                                                                               disable_web_page_preview=True))
                    elif regex.compile('^/dl(@scdlbot)?(\s+.*)?$').match(gottext.strip()):
                        sendaudio(gottext.split()[1],chat_id)
                    elif regex.compile('^http(s)?://(m.)?soundcloud.com/([\w\d_-]+)/([\w\d_-]+)$').match(gottext.strip()):
                        sendaudio(gottext,chat_id)
                else:
                    print('got idkwhat')
                msgfrom = message.get('from')
                print('  from ' + msgfrom['first_name'] + ' ' + msgfrom.get('last_name', '') +
                          ' @' + msgfrom.get('username', ''))

    except (KeyboardInterrupt, SystemExit):
        raise

    except:
        print('smth went wrong')
        #no key 'result' in dictionary -> very bad
        #if r.status_code == requests.codes.ok:
        #if r.json():
        #if r.json()['ok']:
        #if len(updates) != 0:
        #if msgfrom

    finally:
        offset += 1
<|endoftext|>


<commit_before># SWNhanesReader.py

import csv
import re
import numpy as np
import SWSettings as s
import os

type_dict = {
	'RIAGENDR' : int,
	s.PCB153_CODE : float,
	'RIDAGEYR' : int,
	'LBX153'   : float,
	'LBX138LA' : float,
	'LBX180LA' : float,
	'LBX028LA' : float,
	'LBX052LA' : float,
	'LBX101LA' : float,
	'LBX118LA' : float,
	'BMXBMI'   : float
}

class SWNhanesReader(object):

	"""docstring for NhanesReader

	A class designed to read all of the NHANES data
	from a specific sampling year
	and store it in a dictionary

	- So it can be easily accessed by whoever wants to work with the NHANES data.
	"""

	# Public API.
	def __init__(self, nhanes_year='2003-2004', diet = False):
		super(SWNhanesReader, self).__init__()
		self.READ_DIETARY_INFO = diet # save a lot of time reading data if False
		self.nhanes_year = nhanes_year
		self.data = self.obtain_data()

	def concentration_for_seqn_for_pcb(self, seqn, pcb):
		#return self.data[seqn][self.get_nhanes_code_for_pcb(pcb)]
		return self.data.get(seqn).get(self.get_nhanes_code_for_pcb(pcb))

	def get_list_of_seqn_for_pcb(self, pcb='PCB-153'): # Default is PCB-153
		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		return [seqn for seqn in self.data if pcb_string in self.data[seqn]];

	def get_list_of_seqn(self):
		return self.data.keys()

	def get_concentration_list_for_pcb(self, pcb):
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)
		nhanes_code = self.get_nhanes_code_for_pcb(pcb)
		return [self.data.get(x).get(nhanes_code) for x in seqn_list]

	def get_median_concentration_for_pcb_for_gender(self, pcb, female):
		# get all median concentrations for a particular PCB
		# for male or female
		# for 11-20, 21-30, 31-40, and so on.
		start_age = 11
		end_age = 90
		age_gap = s.DEFAULT_AGE_SPREAD

		ages = []
		concentrations = []

		for i in range (start_age, end_age, age_gap):
			age_group_start = i
			age_group_end = i + age_gap - 1
			age, c = self.get_median_concentration_for_pcb_for_age_group_for_gender(pcb, age_group_start, age_group_end, female)
			ages.append(age)
			concentrations.append(c)

		return [ages, concentrations]

	def get_median_concentration_for_all_ages_for_pcb_for_gender(self, pcb, female):

		# Return the median of the reported NHANES PCB concentrations for a specific gender and specific PCB.

		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)

		c_list = []

		gender = 2 if female else 1

		for seqn in seqn_list:
			if self.data[seqn][s.GENDER_CODE] == gender:
				c_list.append(self.data[seqn][pcb_string])

		median = np.median(c_list)

		return median

	def get_median_concentration_for_pcb_for_age_group_for_gender(self, pcb, min_age, max_age, female):
		# get median concentration of a PCB
		# for a certain age group
		# and for male or female.
		ages, concentrations = self.get_concentrations_for_pcb_for_age_group_for_gender(pcb, min_age, max_age, female)
		median_concentration = np.median(concentrations)
		median_age = np.median(ages)
		return [median_age, median_concentration]

	def get_concentrations_for_pcb_for_age_group_for_gender(self, pcb, min_age, max_age, female):
		# return all concentrations of a certain PCB
		# for a certain age group
		# and for male or female
		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)
		gender = self.gender_number_if_female(female)

		ages = []
		concentrations = []

		for seqn in seqn_list:
			if self.data[seqn][s.GENDER_CODE] == gender:
				if self.data[seqn][s.AGE_CODE] >= min_age and self.data[seqn][s.AGE_CODE] <= max_age:
					ages.append(self.data[seqn][s.AGE_CODE])
					concentrations.append(self.data[seqn][pcb_string])

		return [ages, concentrations]

	def get_age_and_concentration_for_pcb_for_gender(self, pcb, female):
		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)
		gender = self.gender_number_if_female(female)

		ages = []
		concentrations = []
		for seqn in seqn_list:
			if self.data[seqn][s.GENDER_CODE] == gender:
				ages.append(self.data[seqn][s.AGE_CODE])
				concentrations.append(self.data[seqn][pcb_string])

		return [ages, concentrations]

	# Private Methods below this line.

	def gender_number_if_female(self, female):
		return 2 if female else 1

	def get_nhanes_code_for_pcb(self, pcb):
		# This method is designed to take any kind of string input (or int maybe)
		# and turn it into the code used in the NHANES data.
		# i.e. Entering 'PCB-153' should give back 'LBX153LA', the lipid adjusted PCB concentration code.
		pcb = str(pcb) # cast to string in case an int was passed.
		pcb_number =  str(re.findall(r'\d+', pcb)[0] if re.findall(r'\d+', pcb) else 0) # extract the PCB congener number
		if len(pcb_number) == 2: pcb_number = '0' + pcb_number # if its 11-99 then a leading 0 must be added
		pcb_string = 'LBX' + pcb_number + 'LA' # construct the string for the code.
		return pcb_string # return

	def cast(self, var):
		return type_dict[var] if var in type_dict else str

	def obtain_data(self):
		ret_dict = {}

		seqn_column = 1
		start_column = 2
		start_row = 1
		header_row = 0

		filenames = [os.path.join(s.NHANES_DATA_PATH, self.folders_dict[self.nhanes_year], name) for name in self.filenames_dict[self.nhanes_year]]

		imported_data = []

		for f in filenames:
			file_import_temp = []
			with open(f, 'rU') as csvfile:
				temp = csv.reader(csvfile, delimiter = ',', quotechar = '"')
				for row in temp:
					file_import_temp.append(row)
			imported_data.append(file_import_temp)

		if imported_data:
			if imported_data[0]:
				for row in imported_data[0][start_row:]:
					seqn = int(row[seqn_column])
					ret_dict.update({seqn : {} })

		for f in imported_data:
			for row in f[start_row:]:
				seqn = int(row[seqn_column])
				for i, header in enumerate(f[header_row][start_column:]):
					header = header.upper()
					column = i + start_column
					try: val = self.cast(header)(row[column])
					except ValueError: continue
					ret_dict[seqn].update({header : val})

		# special read in for diet info
		if self.READ_DIETARY_INFO:
			f = os.path.join(s.NHANES_DATA_PATH, self.folders_dict[self.nhanes_year], self.diet_filename_dict[self.nhanes_year])

			# f = [os.path.join(s.NHANES_DATA_PATH, self.folders_dict[self.nhanes_year], fn) for fn in diet_filename_dict[self.nhanes_year]]

			food_number_column = 2
			raw = []

			for seqn in ret_dict:
				ret_dict[seqn].update({'food_index' : {} })

			with open(f, 'rU') as csvfile:
				temp = csv.reader(csvfile, delimiter = ',', quotechar = '"')
				for row in temp:
					raw.append(row)
			for row in raw[start_row:]:
				seqn = int(row[seqn_column])
				food_index = int(row[food_number_column])
				ret_dict[seqn]['food_index'].update({food_index : {}})
				for i, header in enumerate(raw[header_row][3:]):
					header = header.upper().replace('DRX', 'DR1')
					column = i + 3
					#print 'brool'
					try: val = self.cast(header)(row[column])
					except ValueError: continue
					ret_dict[seqn]['food_index'][food_index].update({header : val})

		return ret_dict

	# Dictionaries containing the appropriate file names.
	filenames_dict = {
		'2003-2004' : ['DEMO_C.csv', 'RHQ_C.csv', 'L28NPB_C.csv', 'BMX_C.csv', 'L28DFP_C.csv', 'DR1TOT_C.csv'],
		'2001-2002' : ['DEMO_B.csv', 'L28POC_B.csv', 'BMX_B.csv', 'DRXTOT_B.csv', 'RHQ_B.csv'],
		'1999-2000' : ['DEMO.csv', 'LAB28POC.csv', 'BMX.csv', 'DRXTOT.csv', 'RHQ.csv']
	}

	diet_filename_dict = {
		'2003-2004' : 'DR1IFF_C.csv',
		'2001-2002' : 'DRXIFF_B.csv',
		'1999-2000' : 'DRXIFF.csv'
 	}

	folders_dict = {
		'2003-2004' : 'NHANES data 2003-2004',
		'2001-2002' : 'NHANES data 2001-2002',
		'1999-2000' : 'NHANES data 1999-2000'
	}
<commit_msg>Quick change<commit_after># SWNhanesReader.py

import csv
import re
import numpy as np
import SWSettings as s
import os

type_dict = {
	'RIAGENDR' : int,
	s.PCB153_CODE : float,
	'RIDAGEYR' : int,
	'LBX153'   : float,
	'LBX138LA' : float,
	'LBX180LA' : float,
	'LBX028LA' : float,
	'LBX052LA' : float,
	'LBX101LA' : float,
	'LBX118LA' : float,
	'BMXBMI'   : float
}

class SWNhanesReader(object):

	READ_DIETARY_INFO = True # Save time if false

	"""docstring for NhanesReader

	A class designed to read all of the NHANES data
	from a specific sampling year
	and store it in a dictionary

	- So it can be easily accessed by whoever wants to work with the NHANES data.
	"""

	# Public API.
	def __init__(self, nhanes_year='2003-2004', diet = False):
		super(SWNhanesReader, self).__init__()
		self.READ_DIETARY_INFO = diet # save a lot of time reading data if False
		self.nhanes_year = nhanes_year
		self.data = self.obtain_data()

	def concentration_for_seqn_for_pcb(self, seqn, pcb):
		#return self.data[seqn][self.get_nhanes_code_for_pcb(pcb)]
		return self.data.get(seqn).get(self.get_nhanes_code_for_pcb(pcb))

	def get_list_of_seqn_for_pcb(self, pcb='PCB-153'): # Default is PCB-153
		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		return [seqn for seqn in self.data if pcb_string in self.data[seqn]];

	def get_list_of_seqn(self):
		return self.data.keys()

	def get_concentration_list_for_pcb(self, pcb):
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)
		nhanes_code = self.get_nhanes_code_for_pcb(pcb)
		return [self.data.get(x).get(nhanes_code) for x in seqn_list]

	def get_median_concentration_for_pcb_for_gender(self, pcb, female):
		# get all median concentrations for a particular PCB
		# for male or female
		# for 11-20, 21-30, 31-40, and so on.
		start_age = 11
		end_age = 90
		age_gap = s.DEFAULT_AGE_SPREAD

		ages = []
		concentrations = []

		for i in range (start_age, end_age, age_gap):
			age_group_start = i
			age_group_end = i + age_gap - 1
			age, c = self.get_median_concentration_for_pcb_for_age_group_for_gender(pcb, age_group_start, age_group_end, female)
			ages.append(age)
			concentrations.append(c)

		return [ages, concentrations]

	def get_median_concentration_for_all_ages_for_pcb_for_gender(self, pcb, female):

		# Return the median of the reported NHANES PCB concentrations for a specific gender and specific PCB.

		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)

		c_list = []

		gender = 2 if female else 1

		for seqn in seqn_list:
			if self.data[seqn][s.GENDER_CODE] == gender:
				c_list.append(self.data[seqn][pcb_string])

		median = np.median(c_list)

		return median

	def get_median_concentration_for_pcb_for_age_group_for_gender(self, pcb, min_age, max_age, female):
		# get median concentration of a PCB
		# for a certain age group
		# and for male or female.
		ages, concentrations = self.get_concentrations_for_pcb_for_age_group_for_gender(pcb, min_age, max_age, female)
		median_concentration = np.median(concentrations)
		median_age = np.median(ages)
		return [median_age, median_concentration]

	def get_concentrations_for_pcb_for_age_group_for_gender(self, pcb, min_age, max_age, female):
		# return all concentrations of a certain PCB
		# for a certain age group
		# and for male or female
		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)
		gender = self.gender_number_if_female(female)

		ages = []
		concentrations = []

		for seqn in seqn_list:
			if self.data[seqn][s.GENDER_CODE] == gender:
				if self.data[seqn][s.AGE_CODE] >= min_age and self.data[seqn][s.AGE_CODE] <= max_age:
					ages.append(self.data[seqn][s.AGE_CODE])
					concentrations.append(self.data[seqn][pcb_string])

		return [ages, concentrations]

	def get_age_and_concentration_for_pcb_for_gender(self, pcb, female):
		pcb_string = self.get_nhanes_code_for_pcb(pcb)
		seqn_list = self.get_list_of_seqn_for_pcb(pcb)
		gender = self.gender_number_if_female(female)

		ages = []
		concentrations = []
		for seqn in seqn_list:
			if self.data[seqn][s.GENDER_CODE] == gender:
				ages.append(self.data[seqn][s.AGE_CODE])
				concentrations.append(self.data[seqn][pcb_string])

		return [ages, concentrations]

	# Private Methods below this line.

	def gender_number_if_female(self, female):
		return 2 if female else 1

	def get_nhanes_code_for_pcb(self, pcb):
		# This method is designed to take any kind of string input (or int maybe)
		# and turn it into the code used in the NHANES data.
		# i.e. Entering 'PCB-153' should give back 'LBX153LA', the lipid adjusted PCB concentration code.
		pcb = str(pcb) # cast to string in case an int was passed.
		pcb_number =  str(re.findall(r'\d+', pcb)[0] if re.findall(r'\d+', pcb) else 0) # extract the PCB congener number
		if len(pcb_number) == 2: pcb_number = '0' + pcb_number # if its 11-99 then a leading 0 must be added
		pcb_string = 'LBX' + pcb_number + 'LA' # construct the string for the code.
		return pcb_string # return

	def cast(self, var):
		return type_dict[var] if var in type_dict else str

	def obtain_data(self):
		ret_dict = {}

		seqn_column = 1
		start_column = 2
		start_row = 1
		header_row = 0

		filenames = [os.path.join(s.NHANES_DATA_PATH, self.folders_dict[self.nhanes_year], name) for name in self.filenames_dict[self.nhanes_year]]

		imported_data = []

		for f in filenames:
			file_import_temp = []
			with open(f, 'rU') as csvfile:
				temp = csv.reader(csvfile, delimiter = ',', quotechar = '"')
				for row in temp:
					file_import_temp.append(row)
			imported_data.append(file_import_temp)

		if imported_data:
			if imported_data[0]:
				for row in imported_data[0][start_row:]:
					seqn = int(row[seqn_column])
					ret_dict.update({seqn : {} })

		for f in imported_data:
			for row in f[start_row:]:
				seqn = int(row[seqn_column])
				for i, header in enumerate(f[header_row][start_column:]):
					header = header.upper()
					column = i + start_column
					try: val = self.cast(header)(row[column])
					except ValueError: continue
					ret_dict[seqn].update({header : val})

		# special read in for diet info
		if self.READ_DIETARY_INFO:
			f = os.path.join(s.NHANES_DATA_PATH, self.folders_dict[self.nhanes_year], self.diet_filename_dict[self.nhanes_year])

			# f = [os.path.join(s.NHANES_DATA_PATH, self.folders_dict[self.nhanes_year], fn) for fn in diet_filename_dict[self.nhanes_year]]

			food_number_column = 2
			raw = []

			for seqn in ret_dict:
				ret_dict[seqn].update({'food_index' : {} })

			with open(f, 'rU') as csvfile:
				temp = csv.reader(csvfile, delimiter = ',', quotechar = '"')
				for row in temp:
					raw.append(row)
			for row in raw[start_row:]:
				seqn = int(row[seqn_column])
				food_index = int(row[food_number_column])
				ret_dict[seqn]['food_index'].update({food_index : {}})
				for i, header in enumerate(raw[header_row][3:]):
					header = header.upper().replace('DRX', 'DR1')
					column = i + 3
					#print 'brool'
					try: val = self.cast(header)(row[column])
					except ValueError: continue
					ret_dict[seqn]['food_index'][food_index].update({header : val})

		return ret_dict

	# Dictionaries containing the appropriate file names.
	filenames_dict = {
		'2003-2004' : ['DEMO_C.csv', 'RHQ_C.csv', 'L28NPB_C.csv', 'BMX_C.csv', 'L28DFP_C.csv', 'DR1TOT_C.csv'],
		'2001-2002' : ['DEMO_B.csv', 'L28POC_B.csv', 'BMX_B.csv', 'DRXTOT_B.csv', 'RHQ_B.csv'],
		'1999-2000' : ['DEMO.csv', 'LAB28POC.csv', 'BMX.csv', 'DRXTOT.csv', 'RHQ.csv']
	}

	diet_filename_dict = {
		'2003-2004' : 'DR1IFF_C.csv',
		'2001-2002' : 'DRXIFF_B.csv',
		'1999-2000' : 'DRXIFF.csv'
 	}

	folders_dict = {
		'2003-2004' : 'NHANES data 2003-2004',
		'2001-2002' : 'NHANES data 2001-2002',
		'1999-2000' : 'NHANES data 1999-2000'
	}
<|endoftext|>


<commit_before>import argparse
import json
import logging
import subprocess
import yaml
import re
import collections
import os
import time
import requests
import urllib

import teuthology
from . import misc
from . import provision
from .config import config
from .lockstatus import get_status

log = logging.getLogger(__name__)
# Don't need to see connection pool INFO messages
logging.getLogger("requests.packages.urllib3.connectionpool").setLevel(
    logging.WARNING)

is_vpm = lambda name: 'vpm' in name

def main(ctx):
    if ctx.verbose:
        teuthology.log.setLevel(logging.DEBUG)

    misc.read_config(ctx)

    ret = 0
    user = ctx.owner
    machines = [misc.canonicalize_hostname(m, user=False)
                for m in ctx.machines]
    machines_to_update = []

    if ctx.targets:
        try:
            with file(ctx.targets) as f:
                g = yaml.safe_load_all(f)
                for new in g:
                    if 'targets' in new:
                        for t in new['targets'].iterkeys():
                            machines.append(t)
        except IOError as e:
            raise argparse.ArgumentTypeError(str(e))

    if ctx.f:
        assert ctx.lock or ctx.unlock, \
            '-f is only supported by --lock and --unlock'
    if machines:
        assert ctx.lock or ctx.unlock or ctx.list or ctx.list_targets \
            or ctx.update, \
            'machines cannot be specified with that operation'
    else:
        assert ctx.num_to_lock or ctx.list or ctx.list_targets or \
            ctx.summary or ctx.brief, \
            'machines must be specified for that operation'
    if ctx.all:
        assert ctx.list or ctx.list_targets or ctx.brief, \
            '--all can only be used with --list, --list-targets, and --brief'
        assert ctx.owner is None, \
            '--all and --owner are mutually exclusive'
        assert not machines, \
            '--all and listing specific machines are incompatible'
    if ctx.num_to_lock:
        assert ctx.machine_type, \
            'must specify machine type to lock'

    if ctx.brief or ctx.list or ctx.list_targets:
        assert ctx.desc is None, '--desc does nothing with --list/--brief'

        if machines:
            statuses = []
            for machine in machines:
                machine = misc.canonicalize_hostname(machine)
                status = get_status(machine)
                if status:
                    statuses.append(status)
                else:
                    log.error("Lockserver doesn't know about machine: %s" %
                              machine)
        else:
            statuses = list_locks()
        vmachines = []

        for vmachine in statuses:
            if vmachine['vm_host']:
                if vmachine['locked']:
                    vmachines.append(vmachine['name'])
        if vmachines:
            # Avoid ssh-keyscans for everybody when listing all machines
            # Listing specific machines will update the keys.
            if machines:
                do_update_keys(vmachines)
                statuses = [get_status(machine)
                            for machine in machines]
            else:
                statuses = list_locks()
        if statuses:
            if ctx.machine_type:
                statuses = [_status for _status in statuses
                            if _status['machine_type'] == ctx.machine_type]
            if not machines and ctx.owner is None and not ctx.all:
                ctx.owner = misc.get_user()
            if ctx.owner is not None:
                statuses = [_status for _status in statuses
                            if _status['locked_by'] == ctx.owner]
            if ctx.status is not None:
                statuses = [_status for _status in statuses
                            if _status['up'] == (ctx.status == 'up')]
            if ctx.locked is not None:
                statuses = [_status for _status in statuses
                            if _status['locked'] == (ctx.locked == 'true')]
            if ctx.desc is not None:
                statuses = [_status for _status in statuses
                            if _status['description'] == ctx.desc]
            if ctx.desc_pattern is not None:
                statuses = [_status for _status in statuses
                            if _status['description'] is not None and
                            _status['description'].find(ctx.desc_pattern) >= 0]

            # When listing, only show the vm_host's name, not every detail
            for s in statuses:
                if not s.get('is_vm', False):
                    continue
                vm_host_name = s.get('vm_host', dict())['name']
                if vm_host_name:
                    s['vm_host'] = vm_host_name
            if ctx.list:
                    print json.dumps(statuses, indent=4)

            elif ctx.brief:
                for s in statuses:
                    locked = "un" if s['locked'] == 0 else "  "
                    mo = re.match('\w+@(\w+?)\..*', s['name'])
                    host = mo.group(1) if mo else s['name']
                    print '{host} {locked}locked {owner} "{desc}"'.format(
                        locked=locked, host=host,
                        owner=s['locked_by'], desc=s['description'])

            else:
                frag = {'targets': {}}
                for f in statuses:
                    frag['targets'][f['name']] = f['ssh_pub_key']
                print yaml.safe_dump(frag, default_flow_style=False)
        else:
            log.error('error retrieving lock statuses')
            ret = 1

    elif ctx.summary:
        do_summary(ctx)
        return 0

    elif ctx.lock:
        for machine in machines:
            if not lock_one(machine, user, ctx.desc):
                ret = 1
                if not ctx.f:
                    return ret
            else:
                machines_to_update.append(machine)
                provision.create_if_vm(ctx, machine)
    elif ctx.unlock:
        if ctx.owner is None and user is None:
            user = misc.get_user()
        # If none of them are vpm, do them all in one shot
        if not filter(is_vpm, machines):
            res = unlock_many(machines, user)
            return 0 if res else 1
        for machine in machines:
            if not unlock_one(ctx, machine, user):
                ret = 1
                if not ctx.f:
                    return ret
            else:
                machines_to_update.append(machine)
    elif ctx.num_to_lock:
        result = lock_many(ctx, ctx.num_to_lock, ctx.machine_type, user,
                           ctx.desc)
        if not result:
            ret = 1
        else:
            machines_to_update = result.keys()
            if ctx.machine_type == 'vps':
                shortnames = ' '.join(
                    [misc.decanonicalize_hostname(name) for name in
                     result.keys()]
                )
                if len(result) < ctx.num_to_lock:
                    log.error("Locking failed.")
                    for machn in result:
                        unlock_one(ctx, machn)
                    ret = 1
                else:
                    log.info("Successfully Locked:\n%s\n" % shortnames)
                    log.info(
                        "Unable to display keys at this time (virtual " +
                        "machines are booting).")
                    log.info(
                        "Please run teuthology-lock --list-targets %s once " +
                        "these machines come up.",
                        shortnames)
            else:
                print yaml.safe_dump(
                    dict(targets=result),
                    default_flow_style=False)
    elif ctx.update:
        assert ctx.desc is not None or ctx.status is not None, \
            'you must specify description or status to update'
        assert ctx.owner is None, 'only description and status may be updated'
        machines_to_update = machines

        if ctx.desc is not None or ctx.status is not None:
            for machine in machines_to_update:
                update_lock(machine, ctx.desc, ctx.status)

    return ret

def lock_many(ctx, num, machinetype, user=None, description=None):
    machinetypes = misc.get_multi_machine_types(machinetype)
    if user is None:
        user = misc.get_user()
    for machinetype in machinetypes:
        uri = os.path.join(config.lock_server, 'nodes', 'lock_many', '')
        data = dict(
            locked_by=user,
            count=num,
            machine_type=machinetype,
            description=description,
        )
        response = requests.post(
            uri,
            data=json.dumps(data),
            headers={'content-type': 'application/json'},
        )
        if response.ok:
            machines = {misc.canonicalize_hostname(machine['name']):
                        machine['ssh_pub_key'] for machine in response.json()}
            log.debug('locked {machines}'.format(
                machines=', '.join(machines.keys())))
            if machinetype == 'vps':
                ok_machs = {}
                for machine in machines:
                    if provision.create_if_vm(ctx, machine):
                        ok_machs[machine] = machines[machine]
                    else:
                        log.error('Unable to create virtual machine: %s',
                                  machine)
                        unlock_one(ctx, machine)
                return ok_machs
            return machines
        elif response.status_code == 503:
            log.error('Insufficient nodes available to lock %d %s nodes.',
                      num, machinetype)
            log.error(response.text)
        else:
            log.error('Could not lock %d %s nodes, reason: unknown.',
                      num, machinetype)
    return []

def lock_one(name, user=None, description=None):
    if user is None:
        user = misc.get_user()
    request = dict(name=name, locked=True, locked_by=user,
                   description=description)
    uri = os.path.join(config.lock_server, 'nodes', name, 'lock', '')
    response = requests.put(uri, json.dumps(request))
    success = response.ok
    if success:
        log.debug('locked %s as %s', name, user)
    else:
        try:
            reason = response.json().get('message')
        except ValueError:
            reason = str(response.status_code)
        log.error('failed to lock {node}. reason: {reason}'.format(
            node=name, reason=reason))
    return response

def unlock_many(names, user):
    uri = os.path.join(config.lock_server, 'nodes', 'unlock_many', '')
    data = dict(
        locked_by=user,
        names=names,
    )
    response = requests.post(
        uri,
        data=json.dumps(data),
        headers={'content-type': 'application/json'},
    )
    if response.ok:
        log.debug("Unlocked: %s", ', '.join(names))
    else:
        log.error("Failed to unlock: %s", ', '.join(names))
    return response.ok

def unlock_one(ctx, name, user=None):
    if user is None:
        user = misc.get_user()
    request = dict(name=name, locked=False, locked_by=user, description=None)
    uri = os.path.join(config.lock_server, 'nodes', name, 'lock', '')
    response = requests.put(uri, json.dumps(request))
    success = response.ok
    if success:
        log.debug('unlocked %s', name)
        if not provision.destroy_if_vm(ctx, name):
            log.error('downburst destroy failed for %s', name)
            log.info('%s is not locked' % name)
    else:
        try:
            reason = response.json().get('message')
        except ValueError:
            reason = str(response.status_code)
        log.error('failed to unlock {node}. reason: {reason}'.format(
            node=name, reason=reason))
    return success

def list_locks(keyed_by_name=False, **kwargs):
    uri = os.path.join(config.lock_server, 'nodes', '')
    if kwargs:
        if 'machine_type' in kwargs:
            kwargs['machine_type'] = kwargs['machine_type'].replace(',','|')
        uri += '?' + urllib.urlencode(kwargs)
    response = requests.get(uri)
    success = response.ok
    if success:
        if not keyed_by_name:
            return response.json()
        else:
            return {node['name']: node
                    for node in response.json()}
    return None

def update_lock(name, description=None, status=None, ssh_pub_key=None):
    status_info = get_status(name)
    if status_info['is_vm']:
        ssh_key = None
        while not ssh_key:
            time.sleep(10)
            ssh_key = ssh_keyscan(name)
    updated = {}
    if description is not None:
        updated['description'] = description
    if status is not None:
        updated['up'] = (status == 'up')
    if ssh_pub_key is not None:
        updated['ssh_pub_key'] = ssh_pub_key

    if updated:
        uri = os.path.join(config.lock_server, 'nodes', name, '')
        response = requests.put(
            uri,
            json.dumps(updated))
        return response.ok
    return True

def ssh_keyscan(hostnames):
    """
    Fetch the SSH public key of one or more hosts
    """
    args = ['ssh-keyscan', '-t', 'rsa']
    if isinstance(hostnames, basestring):
        args.append(hostnames)
    else:
        args.extend(hostnames)
    p = subprocess.Popen(
        args=args,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    p.wait()

    keys_dict = dict()
    for line in p.stdout.readlines():
        host, key = line.strip().split(' ', 1)
        keys_dict[host] = key
    return keys_dict

def updatekeys(ctx):
    loglevel = logging.INFO
    if ctx.verbose:
        loglevel = logging.DEBUG

    logging.basicConfig(
        level=loglevel,
    )

    misc.read_config(ctx)

    machines = [misc.canonicalize_hostname(m) for m in ctx.machines]

    if ctx.targets:
        try:
            with file(ctx.targets) as f:
                g = yaml.safe_load_all(f)
                for new in g:
                    if 'targets' in new:
                        for t in new['targets'].iterkeys():
                            machines.append(t)
        except IOError as e:
            raise argparse.ArgumentTypeError(str(e))

    return do_update_keys(machines)

def do_update_keys(machines):
    reference = list_locks(keyed_by_name=True, up=True)
    if not machines:
        machines = reference.keys()
    keys_dict = ssh_keyscan(machines)
    return push_new_keys(keys_dict, reference)

def push_new_keys(keys_dict, reference):
    ret = 0
    for hostname, pubkey in keys_dict.iteritems():
        log.info('Checking %s', hostname)
        if reference[hostname]['ssh_pub_key'] != pubkey:
            log.info('New key found. Updating...')
            if not update_lock(hostname, ssh_pub_key=pubkey):
                log.error('failed to update %s!', hostname)
                ret = 1
    return ret

def do_summary(ctx):
    lockd = collections.defaultdict(lambda: [0, 0, 'unknown'])
    if ctx.machine_type:
        locks = list_locks(machine_type=ctx.machine_type)
    else:
        locks = list_locks()
    for l in locks:
        who = l['locked_by'] if l['locked'] == 1 \
            else '(free)', l['machine_type']
        lockd[who][0] += 1
        lockd[who][1] += 1 if l['up'] else 0
        lockd[who][2] = l['machine_type']

    locks = sorted([p for p in lockd.iteritems()
                    ], key=lambda sort: (sort[1][2], sort[1][0]))
    total_count, total_up = 0, 0
    print "TYPE     COUNT  UP  OWNER"

    for (owner, (count, upcount, machinetype)) in locks:
            # if machinetype == spectype:
            print "{machinetype:8s} {count:3d}  {up:3d}  {owner}".format(
                count=count, up=upcount, owner=owner[0],
                machinetype=machinetype)
            total_count += count
            total_up += upcount

    print "         ---  ---"
    print "{cnt:12d}  {up:3d}".format(cnt=total_count, up=total_up)
<commit_msg>Hostname handling for updatekeys<commit_after>import argparse
import json
import logging
import subprocess
import yaml
import re
import collections
import os
import time
import requests
import urllib

import teuthology
from . import misc
from . import provision
from .config import config
from .lockstatus import get_status

log = logging.getLogger(__name__)
# Don't need to see connection pool INFO messages
logging.getLogger("requests.packages.urllib3.connectionpool").setLevel(
    logging.WARNING)

is_vpm = lambda name: 'vpm' in name

def main(ctx):
    if ctx.verbose:
        teuthology.log.setLevel(logging.DEBUG)

    misc.read_config(ctx)

    ret = 0
    user = ctx.owner
    machines = [misc.canonicalize_hostname(m, user=False)
                for m in ctx.machines]
    machines_to_update = []

    if ctx.targets:
        try:
            with file(ctx.targets) as f:
                g = yaml.safe_load_all(f)
                for new in g:
                    if 'targets' in new:
                        for t in new['targets'].iterkeys():
                            machines.append(t)
        except IOError as e:
            raise argparse.ArgumentTypeError(str(e))

    if ctx.f:
        assert ctx.lock or ctx.unlock, \
            '-f is only supported by --lock and --unlock'
    if machines:
        assert ctx.lock or ctx.unlock or ctx.list or ctx.list_targets \
            or ctx.update, \
            'machines cannot be specified with that operation'
    else:
        assert ctx.num_to_lock or ctx.list or ctx.list_targets or \
            ctx.summary or ctx.brief, \
            'machines must be specified for that operation'
    if ctx.all:
        assert ctx.list or ctx.list_targets or ctx.brief, \
            '--all can only be used with --list, --list-targets, and --brief'
        assert ctx.owner is None, \
            '--all and --owner are mutually exclusive'
        assert not machines, \
            '--all and listing specific machines are incompatible'
    if ctx.num_to_lock:
        assert ctx.machine_type, \
            'must specify machine type to lock'

    if ctx.brief or ctx.list or ctx.list_targets:
        assert ctx.desc is None, '--desc does nothing with --list/--brief'

        if machines:
            statuses = []
            for machine in machines:
                machine = misc.canonicalize_hostname(machine)
                status = get_status(machine)
                if status:
                    statuses.append(status)
                else:
                    log.error("Lockserver doesn't know about machine: %s" %
                              machine)
        else:
            statuses = list_locks()
        vmachines = []

        for vmachine in statuses:
            if vmachine['vm_host']:
                if vmachine['locked']:
                    vmachines.append(vmachine['name'])
        if vmachines:
            # Avoid ssh-keyscans for everybody when listing all machines
            # Listing specific machines will update the keys.
            if machines:
                do_update_keys(vmachines)
                statuses = [get_status(machine)
                            for machine in machines]
            else:
                statuses = list_locks()
        if statuses:
            if ctx.machine_type:
                statuses = [_status for _status in statuses
                            if _status['machine_type'] == ctx.machine_type]
            if not machines and ctx.owner is None and not ctx.all:
                ctx.owner = misc.get_user()
            if ctx.owner is not None:
                statuses = [_status for _status in statuses
                            if _status['locked_by'] == ctx.owner]
            if ctx.status is not None:
                statuses = [_status for _status in statuses
                            if _status['up'] == (ctx.status == 'up')]
            if ctx.locked is not None:
                statuses = [_status for _status in statuses
                            if _status['locked'] == (ctx.locked == 'true')]
            if ctx.desc is not None:
                statuses = [_status for _status in statuses
                            if _status['description'] == ctx.desc]
            if ctx.desc_pattern is not None:
                statuses = [_status for _status in statuses
                            if _status['description'] is not None and
                            _status['description'].find(ctx.desc_pattern) >= 0]

            # When listing, only show the vm_host's name, not every detail
            for s in statuses:
                if not s.get('is_vm', False):
                    continue
                vm_host_name = s.get('vm_host', dict())['name']
                if vm_host_name:
                    s['vm_host'] = vm_host_name
            if ctx.list:
                    print json.dumps(statuses, indent=4)

            elif ctx.brief:
                for s in statuses:
                    locked = "un" if s['locked'] == 0 else "  "
                    mo = re.match('\w+@(\w+?)\..*', s['name'])
                    host = mo.group(1) if mo else s['name']
                    print '{host} {locked}locked {owner} "{desc}"'.format(
                        locked=locked, host=host,
                        owner=s['locked_by'], desc=s['description'])

            else:
                frag = {'targets': {}}
                for f in statuses:
                    frag['targets'][f['name']] = f['ssh_pub_key']
                print yaml.safe_dump(frag, default_flow_style=False)
        else:
            log.error('error retrieving lock statuses')
            ret = 1

    elif ctx.summary:
        do_summary(ctx)
        return 0

    elif ctx.lock:
        for machine in machines:
            if not lock_one(machine, user, ctx.desc):
                ret = 1
                if not ctx.f:
                    return ret
            else:
                machines_to_update.append(machine)
                provision.create_if_vm(ctx, machine)
    elif ctx.unlock:
        if ctx.owner is None and user is None:
            user = misc.get_user()
        # If none of them are vpm, do them all in one shot
        if not filter(is_vpm, machines):
            res = unlock_many(machines, user)
            return 0 if res else 1
        for machine in machines:
            if not unlock_one(ctx, machine, user):
                ret = 1
                if not ctx.f:
                    return ret
            else:
                machines_to_update.append(machine)
    elif ctx.num_to_lock:
        result = lock_many(ctx, ctx.num_to_lock, ctx.machine_type, user,
                           ctx.desc)
        if not result:
            ret = 1
        else:
            machines_to_update = result.keys()
            if ctx.machine_type == 'vps':
                shortnames = ' '.join(
                    [misc.decanonicalize_hostname(name) for name in
                     result.keys()]
                )
                if len(result) < ctx.num_to_lock:
                    log.error("Locking failed.")
                    for machn in result:
                        unlock_one(ctx, machn)
                    ret = 1
                else:
                    log.info("Successfully Locked:\n%s\n" % shortnames)
                    log.info(
                        "Unable to display keys at this time (virtual " +
                        "machines are booting).")
                    log.info(
                        "Please run teuthology-lock --list-targets %s once " +
                        "these machines come up.",
                        shortnames)
            else:
                print yaml.safe_dump(
                    dict(targets=result),
                    default_flow_style=False)
    elif ctx.update:
        assert ctx.desc is not None or ctx.status is not None, \
            'you must specify description or status to update'
        assert ctx.owner is None, 'only description and status may be updated'
        machines_to_update = machines

        if ctx.desc is not None or ctx.status is not None:
            for machine in machines_to_update:
                update_lock(machine, ctx.desc, ctx.status)

    return ret

def lock_many(ctx, num, machinetype, user=None, description=None):
    machinetypes = misc.get_multi_machine_types(machinetype)
    if user is None:
        user = misc.get_user()
    for machinetype in machinetypes:
        uri = os.path.join(config.lock_server, 'nodes', 'lock_many', '')
        data = dict(
            locked_by=user,
            count=num,
            machine_type=machinetype,
            description=description,
        )
        response = requests.post(
            uri,
            data=json.dumps(data),
            headers={'content-type': 'application/json'},
        )
        if response.ok:
            machines = {misc.canonicalize_hostname(machine['name']):
                        machine['ssh_pub_key'] for machine in response.json()}
            log.debug('locked {machines}'.format(
                machines=', '.join(machines.keys())))
            if machinetype == 'vps':
                ok_machs = {}
                for machine in machines:
                    if provision.create_if_vm(ctx, machine):
                        ok_machs[machine] = machines[machine]
                    else:
                        log.error('Unable to create virtual machine: %s',
                                  machine)
                        unlock_one(ctx, machine)
                return ok_machs
            return machines
        elif response.status_code == 503:
            log.error('Insufficient nodes available to lock %d %s nodes.',
                      num, machinetype)
            log.error(response.text)
        else:
            log.error('Could not lock %d %s nodes, reason: unknown.',
                      num, machinetype)
    return []

def lock_one(name, user=None, description=None):
    if user is None:
        user = misc.get_user()
    request = dict(name=name, locked=True, locked_by=user,
                   description=description)
    uri = os.path.join(config.lock_server, 'nodes', name, 'lock', '')
    response = requests.put(uri, json.dumps(request))
    success = response.ok
    if success:
        log.debug('locked %s as %s', name, user)
    else:
        try:
            reason = response.json().get('message')
        except ValueError:
            reason = str(response.status_code)
        log.error('failed to lock {node}. reason: {reason}'.format(
            node=name, reason=reason))
    return response

def unlock_many(names, user):
    uri = os.path.join(config.lock_server, 'nodes', 'unlock_many', '')
    data = dict(
        locked_by=user,
        names=names,
    )
    response = requests.post(
        uri,
        data=json.dumps(data),
        headers={'content-type': 'application/json'},
    )
    if response.ok:
        log.debug("Unlocked: %s", ', '.join(names))
    else:
        log.error("Failed to unlock: %s", ', '.join(names))
    return response.ok

def unlock_one(ctx, name, user=None):
    if user is None:
        user = misc.get_user()
    request = dict(name=name, locked=False, locked_by=user, description=None)
    uri = os.path.join(config.lock_server, 'nodes', name, 'lock', '')
    response = requests.put(uri, json.dumps(request))
    success = response.ok
    if success:
        log.debug('unlocked %s', name)
        if not provision.destroy_if_vm(ctx, name):
            log.error('downburst destroy failed for %s', name)
            log.info('%s is not locked' % name)
    else:
        try:
            reason = response.json().get('message')
        except ValueError:
            reason = str(response.status_code)
        log.error('failed to unlock {node}. reason: {reason}'.format(
            node=name, reason=reason))
    return success

def list_locks(keyed_by_name=False, **kwargs):
    uri = os.path.join(config.lock_server, 'nodes', '')
    if kwargs:
        if 'machine_type' in kwargs:
            kwargs['machine_type'] = kwargs['machine_type'].replace(',','|')
        uri += '?' + urllib.urlencode(kwargs)
    response = requests.get(uri)
    success = response.ok
    if success:
        if not keyed_by_name:
            return response.json()
        else:
            return {node['name']: node
                    for node in response.json()}
    return None

def update_lock(name, description=None, status=None, ssh_pub_key=None):
    status_info = get_status(name)
    if status_info['is_vm']:
        ssh_key = None
        while not ssh_key:
            time.sleep(10)
            ssh_key = ssh_keyscan(name)
    updated = {}
    if description is not None:
        updated['description'] = description
    if status is not None:
        updated['up'] = (status == 'up')
    if ssh_pub_key is not None:
        updated['ssh_pub_key'] = ssh_pub_key

    if updated:
        uri = os.path.join(config.lock_server, 'nodes', name, '')
        response = requests.put(
            uri,
            json.dumps(updated))
        return response.ok
    return True

def ssh_keyscan(hostnames):
    """
    Fetch the SSH public key of one or more hosts
    """
    args = ['ssh-keyscan', '-t', 'rsa']
    if isinstance(hostnames, basestring):
        args.append(hostnames)
    else:
        args.extend(hostnames)
    p = subprocess.Popen(
        args=args,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    p.wait()

    keys_dict = dict()
    for line in p.stdout.readlines():
        host, key = line.strip().split(' ', 1)
        keys_dict[host] = key
    return keys_dict

def updatekeys(ctx):
    loglevel = logging.INFO
    if ctx.verbose:
        loglevel = logging.DEBUG

    logging.basicConfig(
        level=loglevel,
    )

    misc.read_config(ctx)

    machines = [misc.canonicalize_hostname(m, user=None) for m in ctx.machines]

    if ctx.targets:
        try:
            with file(ctx.targets) as f:
                g = yaml.safe_load_all(f)
                for new in g:
                    if 'targets' in new:
                        for t in new['targets'].iterkeys():
                            machines.append(t)
        except IOError as e:
            raise argparse.ArgumentTypeError(str(e))

    return do_update_keys(machines)

def do_update_keys(machines):
    reference = list_locks(keyed_by_name=True, up=True)
    if not machines:
        machines = reference.keys()
    keys_dict = ssh_keyscan(machines)
    return push_new_keys(keys_dict, reference)

def push_new_keys(keys_dict, reference):
    ret = 0
    for hostname, pubkey in keys_dict.iteritems():
        log.info('Checking %s', hostname)
        if reference[hostname]['ssh_pub_key'] != pubkey:
            log.info('New key found. Updating...')
            if not update_lock(hostname, ssh_pub_key=pubkey):
                log.error('failed to update %s!', hostname)
                ret = 1
    return ret

def do_summary(ctx):
    lockd = collections.defaultdict(lambda: [0, 0, 'unknown'])
    if ctx.machine_type:
        locks = list_locks(machine_type=ctx.machine_type)
    else:
        locks = list_locks()
    for l in locks:
        who = l['locked_by'] if l['locked'] == 1 \
            else '(free)', l['machine_type']
        lockd[who][0] += 1
        lockd[who][1] += 1 if l['up'] else 0
        lockd[who][2] = l['machine_type']

    locks = sorted([p for p in lockd.iteritems()
                    ], key=lambda sort: (sort[1][2], sort[1][0]))
    total_count, total_up = 0, 0
    print "TYPE     COUNT  UP  OWNER"

    for (owner, (count, upcount, machinetype)) in locks:
            # if machinetype == spectype:
            print "{machinetype:8s} {count:3d}  {up:3d}  {owner}".format(
                count=count, up=upcount, owner=owner[0],
                machinetype=machinetype)
            total_count += count
            total_up += upcount

    print "         ---  ---"
    print "{cnt:12d}  {up:3d}".format(cnt=total_count, up=total_up)
<|endoftext|>


<commit_before>import distutils.core
import errno
import logging
import os
import re
import shutil
import subprocess

log = logging.getLogger("flambe")

data_dir = os.path.join(os.path.dirname(__file__), "data")
cache_dir = ".flambe-cache"

class FlambeException (Exception): pass

def new (path):
    shutil.copytree(data_dir+"/scaffold", path)

def run (config, platform, debug=False):
    build(config, [platform], debug)
    id = get(config, "id")
    if platform == "android":
        log.info("")
        apk = "build/main-android.apk"
        log.info("Installing: " + apk)
        adt(["-uninstallApp", "-platform", "android", "-appid", id], verbose=False, output=False, check=False)
        adt(["-installApp", "-platform", "android", "-package", apk], verbose=False)
        adt(["-launchApp", "-platform", "android", "-appid", id], verbose=False)
        if debug:
            log.info("")
            # Clear the log, then start tailing it
            adb(["logcat", "-c"], verbose=False)
            adb(["logcat", "-v", "raw", "-s", "air.%s:V" % id], verbose=False)

def build (config, platforms=[], debug=False):
    common_flags = ["-lib", "flambe", "-cp", "src"]

    common_flags += ["-main", get(config, "main")]
    common_flags += ["-dce", "full"]
    if debug: common_flags += ["-debug", "--no-opt", "--no-inline"]
    else: common_flags += ["--no-traces"]

    common_flags += to_list(get(config, "haxe_flags", []))

    mkdir_p(cache_dir)
    mkdir_p("build/web/targets")
    distutils.dir_util.copy_tree("web", "build/web")
    shutil.copy(data_dir+"/flambe.js", "build/web")

    shutil.rmtree("build/web/assets", ignore_errors=True)
    shutil.copytree("assets", "build/web/assets")

    # Flags common to all swf-based targets (flash, android, ios)
    swf_flags = ["--flash-strict", "-D", "native_trace",
        "-swf-header", "640:480:60:000000", "-lib", "hxsl"]
    if debug: swf_flags += ["-D", "fdb", "-D", "advanced-telemetry"]

    def build_html ():
        html_flags = ["-D", "html"]
        unminified = cache_dir+"/main-html.unminified.js"
        js = "build/web/targets/main-html.js"
        log.info("Building: " + js)
        if debug:
            haxe(common_flags + html_flags + ["-js", js])
        else:
            # Minify release builds
            haxe(common_flags + html_flags + ["-js", unminified])
            minify_js([unminified], js, strict=True)

    def build_flash ():
        swf = "build/web/targets/main-flash.swf"
        flash_flags = swf_flags + ["-swf-version", "11", "-swf", swf]
        log.info("Building: " + swf)
        haxe(common_flags + flash_flags)

    def build_air (flags):
        # Prepare the assets directory
        mkdir_p(cache_dir+"/air")
        shutil.rmtree(cache_dir+"/air/assets", ignore_errors=True)
        shutil.copytree("assets", cache_dir+"/air/assets",
            ignore=shutil.ignore_patterns("*.ogg", "*.wav", "*.m4a"))

        # Build the swf
        air_flags = swf_flags + ["-lib", "air3", "-swf-version", "11.2", "-D", "flambe_air"]
        haxe(common_flags + air_flags + flags)

    def generate_air_xml (swf, output):
        from xml.dom.minidom import parseString
        from textwrap import dedent
        doc = parseString(dedent("""
            <application xmlns="http://ns.adobe.com/air/application/3.7">
              <id>"""+get(config, "id")+"""</id>
              <versionNumber>"""+get(config, "version")+"""</versionNumber>
              <filename>"""+get(config, "name")+"""</filename>
              <initialWindow>
                <content>"""+swf+"""</content>
                <renderMode>direct</renderMode>
              </initialWindow>
            </application>"""))
        icons = doc.createElement("icon")
        for icon in os.listdir("icons"):
            # Only include properly named square icons
            if re.match("(\d+)x\\1.png", icon):
                size = os.path.splitext(icon)[0]
                image = doc.createElement("image"+size)
                image.appendChild(doc.createTextNode("icons/"+icon))
                icons.appendChild(image)
        doc.firstChild.appendChild(icons)
        with open(output, "w") as file:
            doc.writexml(file)

    def build_android ():
        apk = "build/main-android.apk"
        log.info("Building: " + apk)

        swf = "main-android.swf"
        build_air(["-swf", cache_dir+"/air/"+swf])

        # Generate a dummy certificate if it doesn't exist
        cert = cache_dir+"/air/certificate-android.p12"
        try:
            with open(cert): pass
        except IOError:
            adt(["-certificate", "-cn", "SelfSign", "-validityPeriod", "25", "2048-RSA", cert, "password"])

        xml = cache_dir+"/air/config-android.xml"
        generate_air_xml(swf, xml)

        apk_type = "apk-debug" if debug else "apk-captive-runtime"
        adt(["-package", "-target", apk_type, "-storetype", "pkcs12",
            "-keystore", cert, "-storepass", "password", apk, xml, "icons",
            "-C", cache_dir+"/air", swf, "assets"])

    def build_ios ():
        ipa = "build/main-ios.ipa"
        log.info("Building: " + ipa)

        swf = "main-ios.swf"
        build_air(["-swf", cache_dir+"/air/"+swf])

        xml = cache_dir+"/air/config-ios.xml"
        generate_air_xml(swf, xml)

        ipa_type = "ipa-debug" if debug else "ipa-ad-hoc" # or maybe ipa-test?
        adt(["-package", "-target", ipa_type,
            "-provisioning-profile", "ios/development.mobileprovision",
            "-storetype", "pkcs12", "-keystore", "ios/development.p12", "-storepass", "password",
            ipa, xml, "icons", "-C", cache_dir+"/air", swf, "assets"])

    builders = {
        "html": build_html,
        "flash": build_flash,
        "ios": build_ios,
        "android": build_android,
    }
    for ii, platform in enumerate(platforms):
        if ii != 0: log.info("")
        builders[platform]()

def clean ():
    shutil.rmtree("build", ignore_errors=True)
    shutil.rmtree(cache_dir, ignore_errors=True)

def haxe (flags, **kwargs):
    run_command(["haxe"] + flags, **kwargs)

def adt (flags, **kwargs):
    run_command(["adt"] + flags, **kwargs)

def adb (flags, **kwargs):
    run_command(["adb"] + flags, **kwargs)

def minify_js (inputs, output, strict=False):
    from textwrap import dedent
    command=["java", "-jar", data_dir+"/closure.jar",
        "--warning_level", "QUIET",
        "--js_output_file", output,
        "--output_wrapper", dedent("""\
            /**
             * Cooked with Flambe
             * https://github.com/aduros/flambe
             */
            %output%""")]
    for input in inputs: command += ["--js", input]
    if strict: command += ["--language_in", "ES5_STRICT"]
    run_command(command, verbose=False)

def run_command (command, verbose=True, output=True, check=True):
    if verbose: log.info(" ".join(command))
    stream = None if output else subprocess.PIPE
    code = subprocess.call(command, stdout=stream, stderr=stream)
    if check and code != 0: raise FlambeException()

def to_list (o):
    if isinstance(o, list): return o
    if isinstance(o, str): return o.split()
    return [o]

def get (config, name, default=None):
    if name in config:
        return config[name]
    elif default is not None:
        return default
    else:
        raise FlambeException("Missing required entry in config file: %s" % name)

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

from twisted.internet import reactor
from twisted.web import static, server
from datetime import datetime, timedelta
class Server ():
    http_port = 5000

    def run (self):
        log.info("Serving on http://localhost:%s" % self.http_port)

        class CacheableFile(static.File):
            def render (self, request):
                # Set perma-cache headers on asset requests
                if "v" in request.args:
                    expires = datetime.utcnow() + timedelta(days=(25 * 365))
                    request.setHeader("Expires", expires.strftime("%a, %d %b %Y %H:%M:%S GMT"))
                    request.setHeader("Cache-Control", "max-age=315360000")
                return super(CacheableFile, self).render(request)

        root = CacheableFile("build/web")
        reactor.listenTCP(self.http_port, server.Site(root))
        reactor.run()
<commit_msg>Handle numeric versions.<commit_after>import distutils.core
import errno
import logging
import os
import re
import shutil
import subprocess

log = logging.getLogger("flambe")

data_dir = os.path.join(os.path.dirname(__file__), "data")
cache_dir = ".flambe-cache"

class FlambeException (Exception): pass

def new (path):
    shutil.copytree(data_dir+"/scaffold", path)

def run (config, platform, debug=False):
    build(config, [platform], debug)
    id = get(config, "id")
    if platform == "android":
        log.info("")
        apk = "build/main-android.apk"
        log.info("Installing: " + apk)
        adt(["-uninstallApp", "-platform", "android", "-appid", id], verbose=False, output=False, check=False)
        adt(["-installApp", "-platform", "android", "-package", apk], verbose=False)
        adt(["-launchApp", "-platform", "android", "-appid", id], verbose=False)
        if debug:
            log.info("")
            # Clear the log, then start tailing it
            adb(["logcat", "-c"], verbose=False)
            adb(["logcat", "-v", "raw", "-s", "air.%s:V" % id], verbose=False)

def build (config, platforms=[], debug=False):
    common_flags = ["-lib", "flambe", "-cp", "src"]

    common_flags += ["-main", get(config, "main")]
    common_flags += ["-dce", "full"]
    if debug: common_flags += ["-debug", "--no-opt", "--no-inline"]
    else: common_flags += ["--no-traces"]

    common_flags += to_list(get(config, "haxe_flags", []))

    mkdir_p(cache_dir)
    mkdir_p("build/web/targets")
    distutils.dir_util.copy_tree("web", "build/web")
    shutil.copy(data_dir+"/flambe.js", "build/web")

    shutil.rmtree("build/web/assets", ignore_errors=True)
    shutil.copytree("assets", "build/web/assets")

    # Flags common to all swf-based targets (flash, android, ios)
    swf_flags = ["--flash-strict", "-D", "native_trace",
        "-swf-header", "640:480:60:000000", "-lib", "hxsl"]
    if debug: swf_flags += ["-D", "fdb", "-D", "advanced-telemetry"]

    def build_html ():
        html_flags = ["-D", "html"]
        unminified = cache_dir+"/main-html.unminified.js"
        js = "build/web/targets/main-html.js"
        log.info("Building: " + js)
        if debug:
            haxe(common_flags + html_flags + ["-js", js])
        else:
            # Minify release builds
            haxe(common_flags + html_flags + ["-js", unminified])
            minify_js([unminified], js, strict=True)

    def build_flash ():
        swf = "build/web/targets/main-flash.swf"
        flash_flags = swf_flags + ["-swf-version", "11", "-swf", swf]
        log.info("Building: " + swf)
        haxe(common_flags + flash_flags)

    def build_air (flags):
        # Prepare the assets directory
        mkdir_p(cache_dir+"/air")
        shutil.rmtree(cache_dir+"/air/assets", ignore_errors=True)
        shutil.copytree("assets", cache_dir+"/air/assets",
            ignore=shutil.ignore_patterns("*.ogg", "*.wav", "*.m4a"))

        # Build the swf
        air_flags = swf_flags + ["-lib", "air3", "-swf-version", "11.2", "-D", "flambe_air"]
        haxe(common_flags + air_flags + flags)

    def generate_air_xml (swf, output):
        from xml.dom.minidom import parseString
        from textwrap import dedent
        doc = parseString(dedent("""
            <application xmlns="http://ns.adobe.com/air/application/3.7">
              <id>"""+get(config, "id")+"""</id>
              <versionNumber>"""+str(get(config, "version"))+"""</versionNumber>
              <filename>"""+get(config, "name")+"""</filename>
              <initialWindow>
                <content>"""+swf+"""</content>
                <renderMode>direct</renderMode>
              </initialWindow>
            </application>"""))
        icons = doc.createElement("icon")
        for icon in os.listdir("icons"):
            # Only include properly named square icons
            if re.match("(\d+)x\\1.png", icon):
                size = os.path.splitext(icon)[0]
                image = doc.createElement("image"+size)
                image.appendChild(doc.createTextNode("icons/"+icon))
                icons.appendChild(image)
        doc.firstChild.appendChild(icons)
        with open(output, "w") as file:
            doc.writexml(file)

    def build_android ():
        apk = "build/main-android.apk"
        log.info("Building: " + apk)

        swf = "main-android.swf"
        build_air(["-swf", cache_dir+"/air/"+swf])

        # Generate a dummy certificate if it doesn't exist
        cert = cache_dir+"/air/certificate-android.p12"
        try:
            with open(cert): pass
        except IOError:
            adt(["-certificate", "-cn", "SelfSign", "-validityPeriod", "25", "2048-RSA", cert, "password"])

        xml = cache_dir+"/air/config-android.xml"
        generate_air_xml(swf, xml)

        apk_type = "apk-debug" if debug else "apk-captive-runtime"
        adt(["-package", "-target", apk_type, "-storetype", "pkcs12",
            "-keystore", cert, "-storepass", "password", apk, xml, "icons",
            "-C", cache_dir+"/air", swf, "assets"])

    def build_ios ():
        ipa = "build/main-ios.ipa"
        log.info("Building: " + ipa)

        swf = "main-ios.swf"
        build_air(["-swf", cache_dir+"/air/"+swf])

        xml = cache_dir+"/air/config-ios.xml"
        generate_air_xml(swf, xml)

        ipa_type = "ipa-debug" if debug else "ipa-ad-hoc" # or maybe ipa-test?
        adt(["-package", "-target", ipa_type,
            "-provisioning-profile", "ios/development.mobileprovision",
            "-storetype", "pkcs12", "-keystore", "ios/development.p12", "-storepass", "password",
            ipa, xml, "icons", "-C", cache_dir+"/air", swf, "assets"])

    builders = {
        "html": build_html,
        "flash": build_flash,
        "ios": build_ios,
        "android": build_android,
    }
    for ii, platform in enumerate(platforms):
        if ii != 0: log.info("")
        builders[platform]()

def clean ():
    shutil.rmtree("build", ignore_errors=True)
    shutil.rmtree(cache_dir, ignore_errors=True)

def haxe (flags, **kwargs):
    run_command(["haxe"] + flags, **kwargs)

def adt (flags, **kwargs):
    run_command(["adt"] + flags, **kwargs)

def adb (flags, **kwargs):
    run_command(["adb"] + flags, **kwargs)

def minify_js (inputs, output, strict=False):
    from textwrap import dedent
    command=["java", "-jar", data_dir+"/closure.jar",
        "--warning_level", "QUIET",
        "--js_output_file", output,
        "--output_wrapper", dedent("""\
            /**
             * Cooked with Flambe
             * https://github.com/aduros/flambe
             */
            %output%""")]
    for input in inputs: command += ["--js", input]
    if strict: command += ["--language_in", "ES5_STRICT"]
    run_command(command, verbose=False)

def run_command (command, verbose=True, output=True, check=True):
    if verbose: log.info(" ".join(command))
    stream = None if output else subprocess.PIPE
    code = subprocess.call(command, stdout=stream, stderr=stream)
    if check and code != 0: raise FlambeException()

def to_list (o):
    if isinstance(o, list): return o
    if isinstance(o, str): return o.split()
    return [o]

def get (config, name, default=None):
    if name in config:
        return config[name]
    elif default is not None:
        return default
    else:
        raise FlambeException("Missing required entry in config file: %s" % name)

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

from twisted.internet import reactor
from twisted.web import static, server
from datetime import datetime, timedelta
class Server ():
    http_port = 5000

    def run (self):
        log.info("Serving on http://localhost:%s" % self.http_port)

        class CacheableFile(static.File):
            def render (self, request):
                # Set perma-cache headers on asset requests
                if "v" in request.args:
                    expires = datetime.utcnow() + timedelta(days=(25 * 365))
                    request.setHeader("Expires", expires.strftime("%a, %d %b %Y %H:%M:%S GMT"))
                    request.setHeader("Cache-Control", "max-age=315360000")
                return super(CacheableFile, self).render(request)

        root = CacheableFile("build/web")
        reactor.listenTCP(self.http_port, server.Site(root))
        reactor.run()
<|endoftext|>