<commit_before>"""Add electron lifetime
"""
import sympy
from sympy.parsing.sympy_parser import parse_expr
from pax import configuration, units
from cax import config
from cax.task import Task

from hax import slow_control

PAX_CONFIG = configuration.load_configuration('XENON1T')

class CorrectionBase(Task):
    "Derive correction"

    def __init__(self):
        self.collection_purity = config.mongo_collection('purity')
        Task.__init__(self)

    def evaluate(self):
        raise NotImplementedError()

    def each_run(self):
        short_key = self.key.split('.')[-1]
        if 'processor' in self.run_doc and \
                        'DEFAULT' in self.run_doc['processor'] and \
                        short_key in self.run_doc['processor']['DEFAULT']:
            return

        if 'end' not in self.run_doc:
            return

        # Fetch the latest electron lifetime fit
        doc = self.collection_purity.find_one(sort=(('calculation_time',
                                                     -1),))

        # Get fit function
        self.function = parse_expr(doc['function'])

        if not config.DATABASE_LOG:
            return

        # Update run database
        self.collection.find_and_modify({'_id': self.run_doc['_id'],
                                         self.key : {'$exists' : False}},
                                        {'$set': {self.key: self.evaluate()}})

class AddElectronLifetime(CorrectionBase):
    """Copy data to here

    If data exists at a reachable host but not here, pull it.
    """
    collection = 'purity'
    key = 'processor.DEFAULT.electron_lifetime_liquid'
    correction_units = units.us

    def evaluate(self):
        # Compute lifetime from this function on this dataset
        lifetime = self.function.evalf(subs={"t" : self.run_doc['start'].timestamp()})
        lifetime = float(lifetime)  # Convert away from Sympy type.

        run_number = self.run_doc['number']
        self.log.info("Run %d: calculated lifetime of %d us" % (run_number,
                                                                lifetime))
        return lifetime * self.correction_units

class AddGains(CorrectionBase):
    """Copy data to here

    If data exists at a reachable host but not here, pull it.
    """
    collection = 'gains'
    key = 'processor.DEFAULT.gains'
    correction_units = units.V # should be 1

    n_channels = 248

    def evaluate(self):
        """Make an array of all PMT gains."""
        n_channels = len(PAX_CONFIG['DEFAULT']['pmts'])
        return [float(self.get_gain(i)) for i in range(n_channels)]

    def get_gain(self, pmt_location):
        """Grab a derived gain.

        pmt_location is the PMT number.  t0 and t1 are datetime objects.

        The variables fed in can be used for making a gain decision.
        """
        # Name of the slow-control variable
        name = slow_control.VARIABLES['pmts']['pmt_%03d_bias_V' % pmt_location]

        # Fetch from slow control
        voltages = slow_control.get_series(name,
                                           time_range=(self.run_doc['start'],
                                                       self.run_doc['end']))

        # If no values found, use default gain.
        if voltages.count() == 0:
            return 2e6

        V = sympy.symbols('V')
        pmt = sympy.symbols('pmt', integer=True)

        result = self.function.evalf(subs={V: voltages.median(),
                                           pmt : pmt_location})
        return float(result) * self.correction_units
<commit_msg>Add gains too<commit_after>"""Add electron lifetime
"""
import sympy
from sympy.parsing.sympy_parser import parse_expr
from pax import configuration, units
from cax import config
from cax.task import Task

from hax import slow_control

PAX_CONFIG = configuration.load_configuration('XENON1T')

class CorrectionBase(Task):
    "Derive correction"

    def __init__(self):
        self.collection_purity = config.mongo_collection('purity')
        Task.__init__(self)

    def evaluate(self):
        raise NotImplementedError()

    def each_run(self):
        short_key = self.key.split('.')[-1]
        if 'processor' in self.run_doc and \
                        'DEFAULT' in self.run_doc['processor'] and \
                        short_key in self.run_doc['processor']['DEFAULT']:
            return

        if 'end' not in self.run_doc:
            return

        # Fetch the latest electron lifetime fit
        doc = self.collection_purity.find_one(sort=(('calculation_time',
                                                     -1),))

        # Get fit function
        self.function = parse_expr(doc['function'])

        if not config.DATABASE_LOG:
            return

        # Update run database
        self.collection.find_and_modify({'_id': self.run_doc['_id'],
                                         self.key : {'$exists' : False}},
                                        {'$set': {self.key: self.evaluate()}})

class AddElectronLifetime(CorrectionBase):
    """Copy data to here

    If data exists at a reachable host but not here, pull it.
    """
    collection = 'purity'
    key = 'processor.DEFAULT.electron_lifetime_liquid'
    correction_units = units.us

    def evaluate(self):
        # Compute lifetime from this function on this dataset
        lifetime = self.function.evalf(subs={"t" : self.run_doc['start'].timestamp()})
        lifetime = float(lifetime)  # Convert away from Sympy type.

        run_number = self.run_doc['number']
        self.log.info("Run %d: calculated lifetime of %d us" % (run_number,
                                                                lifetime))
        return lifetime * self.correction_units

class AddGains(CorrectionBase):
    """Copy data to here

    If data exists at a reachable host but not here, pull it.
    """
    collection = 'gains'
    key = 'processor.DEFAULT.gains'
    correction_units = units.V # should be 1

    n_channels = 248

    def evaluate(self):
        """Make an array of all PMT gains."""
        n_channels = len(PAX_CONFIG['DEFAULT']['pmts'])
        return [float(self.get_gain(i)) for i in range(n_channels)]

    def get_gain(self, pmt_location):
        """Grab a derived gain.

        pmt_location is the PMT number.  t0 and t1 are datetime objects.

        The variables fed in can be used for making a gain decision.
        """
        # Name of the slow-control variable
        name = slow_control.VARIABLES['pmts']['pmt_%03d_bias_V' % pmt_location]

        # Fetch from slow control
        voltages = slow_control.get_series(name,
                                           time_range=(self.run_doc['start'],
                                                       self.run_doc['end']))

        # If no values found, use default gain.
        if voltages.count() == 0:
            return 2e6

        V = sympy.symbols('V')
        pmt = sympy.symbols('pmt', integer=True)
        self.info(voltages.median())
        result = self.function.evalf(subs={V: voltages.median(),
                                           pmt : pmt_location})
        self.info(result)
        return float(result) * self.correction_units
<|endoftext|>


<commit_before>
import sys
import requests

import imaging
import analysis

import serial
import time

ser = serial.Serial(sys.argv[2], 9600)

session_name = sys.argv[1]
print('Sending photo to session: -'+session_name+'-')
session_status_request_error_counter = 0

def looping():
    global session_status_request_error_counter
    sessionstatusrequest_successful = False
    while sessionstatusrequest_successful == False:
        try:
            req = requests.get('http://digitalfoodone.appspot.com/controllersessionstatus', params={'session_name':session_name})
            print(str(req.status_code))
            data = req.json()
            print(data['state'])

            def imagestuff():
                resp = imaging.capture(session_name,'capturing')
                print('Received http status code: '+str(resp.status_code))
                data = resp.json()
                print('Status: '+data['status'])
                analysis_resp = analysis.red_spot(session_name)
                print('Received http status code: '+str(analysis_resp.status_code))
                analysis_data = analysis_resp.json()
                print('Analysis status: '+analysis_data['status'])

            if data['state'] == 'rewind':
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                imagestuff()
            if data['state'] == 'play_one':
                ser.write(chr(70))
                imagestuff()
            if data['state'] == 'play_two':
                print('waiting')
            sessionstatusrequest_successful = True
        except requests.exceptions.Timeout:
            print("session status request: request timeout")
        except requests.exceptions.TooManyRedirects:
            print('session status request: too many redirects')
        except requests.exceptions.RequestException as e:
            print('session status request: request exception:')
            print e
        print('apparently not successful - trying again in 10 sec')
        session_status_request_error_counter += 1
        print('number of status request errors: '+session_status_request_error_counter)
        time.sleep(10)

while True:
    looping()
    time.sleep(8)
<commit_msg>added exception handling to all requests<commit_after>
import sys
import requests

import imaging
import analysis

import serial
import time

ser = serial.Serial(sys.argv[2], 9600)

session_name = sys.argv[1]
print('Sending photo to session: -'+session_name+'-')
session_status_request_error_counter = 0

def looping():
    global session_status_request_error_counter
    sessionstatusrequest_successful = False
    while sessionstatusrequest_successful == False:
        try:
            req = requests.get('http://digitalfoodone.appspot.com/controllersessionstatus', params={'session_name':session_name})
            print(str(req.status_code))
            data = req.json()
            print(data['state'])

            def imagestuff():
                resp = imaging.capture(session_name,'capturing')
                print('Received http status code: '+str(resp.status_code))
                data = resp.json()
                print('Status: '+data['status'])
                analysis_resp = analysis.red_spot(session_name)
                print('Received http status code: '+str(analysis_resp.status_code))
                analysis_data = analysis_resp.json()
                print('Analysis status: '+analysis_data['status'])

            if data['state'] == 'rewind':
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                ser.write(chr(1))
                imagestuff()
            if data['state'] == 'play_one':
                ser.write(chr(70))
                imagestuff()
            if data['state'] == 'play_two':
                print('waiting')
            sessionstatusrequest_successful = True
        except requests.exceptions.Timeout:
            print("session status request: request timeout")
        except requests.exceptions.TooManyRedirects:
            print('session status request: too many redirects')
        except requests.exceptions.RequestException as e:
            print('session status request: request exception:')
            print e
        print('apparently not successful - trying again in 10 sec')
        session_status_request_error_counter += 1
        print('number of status request errors: '+str(session_status_request_error_counter))
        time.sleep(10)

while True:
    looping()
    time.sleep(8)
<|endoftext|>


<commit_before>#! /usr/bin/env python
#
# Python script for the Interoute Virtual Data Centre API:
#   Name: vm_deploy_chooser.py:
#   Purpose: Chooser-based command line interface to deploy a virtual machine
#   Requires: class VDCApiCall in the file vdc_api_call.py
# See the repo: https://github.com/Interoute/API-fun-and-education
#
# VERSION 2
#
# You can pass options via the command line: type 'python vm_deploy_chooser.py -h'
# for usage information
#
# Copyright (C) Interoute Communications Limited, 2015
#
# Original code by Kelcey Damage, 2012 & Kraig Amador, 2012
# Revised by Sandy Walker (Interoute), 2013

from __future__ import print_function
from collections import OrderedDict
import base64
import vdc_api_call as vdc
import getpass
import json
import os
import sys
import pprint
import time
import textwrap
import argparse
import urllib
import urllib2
import hashlib
import hmac

# Print a list and return choice of item. Source: http://stackoverflow.com/questions/13354460/simple-terminal-file-chooser-in-python-libraries
def choose_item_from_list(itemlist, prompt='Please select an item by its number'):
    prompt = prompt + ": "
    # this print is for preliminary message..
    # print(textwrap.dedent() )
    item_dict = {index: value for index, value in enumerate(itemlist)}
    item_display_list = ["%2d. %s" % (key, item_dict[key]) for key in item_dict]
    column_print(item_display_list)
    print('')
    response = int(input(prompt))
    return {'itemindex': response, 'itemcontent': item_dict[response]}

# Print a list by columns. Source: 'col_print' at http://stackoverflow.com/questions/1524126/how-to-print-a-list-more-nicely
# Modified with 'min_length_for_columns' so that short lists must appear in one column
def column_print(lines, term_width=120, indent=0, pad=2, min_length_for_columns=15):
   n_lines = len(lines)
   if n_lines == 0:
      return
   col_width = max(len(line) for line in lines)
   if n_lines < min_length_for_columns:
      n_cols = 1
      col_len = int(n_lines/n_cols) + (0 if n_lines % n_cols == 0 else 1)
   else:
      n_cols = int((term_width + pad - indent)/(col_width + pad))
      n_cols = min(n_lines, max(1, n_cols))
      col_len = int(n_lines/n_cols) + (0 if n_lines % n_cols == 0 else 1)
      if (n_cols - 1) * col_len >= n_lines:
         n_cols -= 1
   cols = [lines[i*col_len : i*col_len + col_len] for i in range(n_cols)]
   rows = list(zip(*cols))
   rows_missed = zip(*[col[len(rows):] for col in cols[:-1]])
   rows.extend(rows_missed)
   for row in rows:
      print(" "*indent + (" "*pad).join(line.ljust(col_width) for line in row))

if __name__ == '__main__':
    # STEP: Parse the command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", default=os.path.join(os.path.expanduser('~'), '.vdcapi'),
                    help="path/name of the config file to be used for the API URL and API keys (default is ~/.vdcapi)")
    parser.add_argument("-m", "--mode", choices=['deploy','d','print','p'], default='print',
                    help="specify the output mode: deploy (deploy the VM) or print (print the API call) (default: print)")
    parser.add_argument("-f", "--format", choices=['json','j','cloudmonkey','c','url'], default='json',
                        help="specify the format of the printed API call: json, cloudmonkey or url (default: json)")
    parser.add_argument("-r", "--region", choices=['Europe','europe','USA','usa','Asia','asia'],
                    default='Europe', help="specify the VDC region: Europe, USA or Asia (default Europe)")
    parser.add_argument("-i", "--iso", action='store_true', help="deploy from an ISO image")
    parser.add_argument("-k", "--keys", action='store_true', help="ask for choice of SSH keys")
    parser.add_argument("-u", "--userdata", action='store_true', help="ask for input of userdata by filename")
    parser.add_argument("-a", "--affinity", action='store_true', help="ask for selection of affinity group(s)")
    vdcRegion = parser.parse_args().region
    config_file = parser.parse_args().config
    mode = parser.parse_args().mode
    printFormat = parser.parse_args().format
    deployFromISO = parser.parse_args().iso
    askForSSHKeys = parser.parse_args().keys
    askForUserdata = parser.parse_args().userdata
    askForAffinityGroup = parser.parse_args().affinity

    # STEP: If config file is found, read its content,
    # else query user for the URL, API key, Secret key
    if os.path.isfile(config_file):
        with open(config_file) as fh:
            data = fh.read()
            config = json.loads(data)
            api_url = config['api_url']
            apiKey = config['api_key']
            secret = config['api_secret']
    else:
        print('API url (e.g. http://10.220.18.115:8080/client/api):', end='')
        api_url = raw_input()
        print('API key:', end='')
        apiKey = raw_input()
        secret = getpass.getpass(prompt='API secret:')

    # STEP: Create the API access object
    api = vdc.VDCApiCall(api_url, apiKey, secret)

#TO ADD: CHOOSE REGION OR FIND ZONES FOR ALL REGIONS THAT THE VDC ACCOUNT CAN ACCESS..........................

    # STEP: Select the zone
    result = api.listZones({'region': vdcRegion})
    zonelist = [zone['name'] for zone in result['zone']]
    zone_ids = [zone['id'] for zone in result['zone']]
    choice = choose_item_from_list(zonelist, prompt="Select the zone?")
    zone_id = zone_ids[choice['itemindex']]
    print("Selected zone: %s, %s\n" % (zone_id, choice['itemcontent']))

    # STEP: Select the template or ISO image
    if deployFromISO:
       print("Error: ISO case not implemented yet")
       exit
    else:
       request = {
        'region': vdcRegion,
        'templatefilter': 'executable',
        'zoneid': zone_id,
       }
       result = api.listTemplates(request)
       templates_sorted = sorted(result['template'], key=lambda item: item['name'].upper())
       templatelist = [template['name'] for template in templates_sorted]
       template_ids = [template['id'] for template in templates_sorted]
       choice = choose_item_from_list(templatelist, prompt="Select the template?")
       template_id = template_ids[choice['itemindex']]
       print("Selected template: %s, '%s'\n" % (template_id, choice['itemcontent']))

    # STEP: Select the compute/service offering
    ramlist = [512,1024,2048,4096,6144,8192,16384,24576,32768,65536,131072]
    choice_cpu = input("Input the number of CPUs? (between 1 and 12): ")
    print('')
    choice_ram = choose_item_from_list(map(lambda x: float(x)/1024, ramlist), prompt="Select the amount of RAM (GBytes)?")
    result = api.listServiceOfferings({'region': vdcRegion})
    serviceoffering_id = [s['id'] for s in result['serviceoffering'] if s['name']=='%d-%d'% (ramlist[choice_ram['itemindex']],choice_cpu)][0]
    print("Selected service offering: %s, \'%s\'\n" % (serviceoffering_id, '%d-%d'% (ramlist[choice_ram['itemindex']],choice_cpu)))

    # (optional) STEP: Select the affinity groups (if any exist)
    if askForAffinityGroup:
       print("Error: Select affinity groups not implemented yet")

    # STEP: Select the network(s)
    result = api.listNetworks({'region': vdcRegion, 'zoneid': zone_id})
    network_ids = [network['id'] for network in result['network']]
    networklist = ['%s (%s)' % (network['displaytext'],network['name']) for network in result['network']]
    network_num = input("Input the number of networks?: ")
    if network_num == 1:
       choice = choose_item_from_list(networklist, prompt="Select the network from the list? (this will be the default)")
       network_id = network_ids[choice['itemindex']]
       print("Selected network: %s, %s\n" % (network_id, choice['itemcontent']))
    else:
       choice = choose_item_from_list(networklist, prompt="Select network 1? (this will be the default)")
       network_id = [network_ids[choice['itemindex']]]
       for i in range(1,network_num):
          choice = choose_item_from_list(networklist, prompt="Select network %d?"%(i+1))
          network_id = network_id + [network_ids[choice['itemindex']]]
       print("Selected networks: %s\n" % (network_id))

    # (optional) STEP: Select keys
    if askForSSHKeys:
       result = api.listSSHKeyPairs({'region': vdcRegion})
       if result=={}:
          print("Error: SSH Keypair asked for, but none are set up in this account in this region")
          askForSSHKeys = False
       else:
          keypairnames = [keypair['name'] for keypair in result['sshkeypair']]
          keypairlist = ['%s (%s)' % (keypair['name'],keypair['fingerprint']) for keypair in result['sshkeypair']]
          choice = choose_item_from_list(keypairlist, prompt="Select SSH keypair?")
          keypairname = keypairnames[choice['itemindex']]

    # (optional) STEP: Input userdata
    if askForUserdata:
       print('')
       userdataFilename = raw_input('Input the path/name for the PLAIN TEXT file with the userdata: ')
       try:
          with open(userdataFilename) as fh:
             userdata=fh.read()
          userdata_b64 = base64.b64encode(userdata)
       except IOError:
          print("Error: Userdata plain text file not found or cannot be opened. Userdata will not be included in the deploy.")
          askForUserdata = False
          pass

    # STEP: Enter the VM 'name' and 'displaytext'
    ## TO DO: (1) create a default VM name, (2) do a test to check new name is unique against existing VMs
    #default_hostname =
    print('')
    hostname = raw_input(
       # 'Input name for the new VM (default %s):' % default_hostname
       'Input name for the new VM (name must be unique): '
    )
    hostname = hostname.strip()
    #if len(hostname) == 0:
    #    hostname = default_hostname

    default_displayname = hostname
    displayname = raw_input(
        'Input displayname for new VM (press ENTER for default: %s):' % default_displayname
    )
    displayname = displayname.strip()
    if len(displayname) == 0:
        displayname = default_displayname

    # STEP: Deploy the VM, or print out the API call
    deploy_params = {
        'region': vdcRegion,
        'serviceofferingid': serviceoffering_id,
        'templateid': template_id,
        'zoneid': zone_id,
        'displayname': displayname,
        'name': hostname,
        'networkids': network_id
    }
    if askForSSHKeys:
        deploy_params['keypair'] = keypairname
    if askForUserdata:
        deploy_params['userdata'] = userdata_b64
    print('')
    if mode=='print':
        print("-------------------------------------------\nRequired deploy command in %s format\n-------------------------------------------\n" % (printFormat))
        if printFormat=='json':
           print(json.dumps(deploy_params))
           print('')
        elif printFormat=='cloudmonkey':
           # For Cloudmonkey, 'region' is set by a separate command
           deploy_params.pop('region')
           params = ["%s=%s" % (key, deploy_params[key]) for key in deploy_params]
           print("NOTE: you need to execute the command 'set region %s' for the following deploy command to work...\n" % (vdcRegion))
           print("deploy virtualmachine " + " ".join(params))
           print('')
        else:
           deploy_params['command'] = 'deployVirtualMachine'
           httprequest = zip(deploy_params.keys(), deploy_params.values())
           httprequest.sort(key=lambda x: x[0].lower())
           httprequest_data = "&".join(["=".join([r[0], urllib.quote_plus(str(r[1]))])
                                for r in httprequest])
           hashStr = "&".join(
               [
                "=".join(
                    [r[0].lower(),
                     str.lower(urllib.quote_plus(str(r[1]))).replace(
                         "+", "%20"
                     )]
                ) for r in httprequest
            ])
           sig = urllib.quote_plus(base64.b64encode(
             hmac.new(
                 secret,
                 hashStr,
                 hashlib.sha1
             ).digest()
           ).strip())
           httprequest_data += "&signature=%s" % sig
           print(api_url + "?" + httprequest_data)
           print('')
    else:
        print("Ready to deploy VM with these parameters:")
        pprint.pprint(json.dumps(deploy_params))
        choice = raw_input("Input D to deploy or any other key to exit:")
        if choice=="d" or choice=="D":
           result = api.deployVirtualMachine(deploy_params)
           job_id = result['jobid']
           pprint.pprint(api.wait_for_job(job_id))
<commit_msg>vm_deploy_chooser.py: Fix bug in url output<commit_after>#! /usr/bin/env python
#
# Python script for the Interoute Virtual Data Centre API:
#   Name: vm_deploy_chooser.py:
#   Purpose: Chooser-based command line interface to deploy a virtual machine
#   Requires: class VDCApiCall in the file vdc_api_call.py
# See the repo: https://github.com/Interoute/API-fun-and-education
#
# VERSION 2
#
# You can pass options via the command line: type 'python vm_deploy_chooser.py -h'
# for usage information
#
# Copyright (C) Interoute Communications Limited, 2015
#
# Original code by Kelcey Damage, 2012 & Kraig Amador, 2012
# Revised by Sandy Walker (Interoute), 2013

from __future__ import print_function
from collections import OrderedDict
import base64
import vdc_api_call as vdc
import getpass
import json
import os
import sys
import pprint
import time
import textwrap
import argparse
import urllib
import urllib2
import hashlib
import hmac

# Print a list and return choice of item. Source: http://stackoverflow.com/questions/13354460/simple-terminal-file-chooser-in-python-libraries
def choose_item_from_list(itemlist, prompt='Please select an item by its number'):
    prompt = prompt + ": "
    # this print is for preliminary message..
    # print(textwrap.dedent() )
    item_dict = {index: value for index, value in enumerate(itemlist)}
    item_display_list = ["%2d. %s" % (key, item_dict[key]) for key in item_dict]
    column_print(item_display_list)
    print('')
    response = int(input(prompt))
    return {'itemindex': response, 'itemcontent': item_dict[response]}

# Print a list by columns. Source: 'col_print' at http://stackoverflow.com/questions/1524126/how-to-print-a-list-more-nicely
# Modified with 'min_length_for_columns' so that short lists must appear in one column
def column_print(lines, term_width=120, indent=0, pad=2, min_length_for_columns=15):
   n_lines = len(lines)
   if n_lines == 0:
      return
   col_width = max(len(line) for line in lines)
   if n_lines < min_length_for_columns:
      n_cols = 1
      col_len = int(n_lines/n_cols) + (0 if n_lines % n_cols == 0 else 1)
   else:
      n_cols = int((term_width + pad - indent)/(col_width + pad))
      n_cols = min(n_lines, max(1, n_cols))
      col_len = int(n_lines/n_cols) + (0 if n_lines % n_cols == 0 else 1)
      if (n_cols - 1) * col_len >= n_lines:
         n_cols -= 1
   cols = [lines[i*col_len : i*col_len + col_len] for i in range(n_cols)]
   rows = list(zip(*cols))
   rows_missed = zip(*[col[len(rows):] for col in cols[:-1]])
   rows.extend(rows_missed)
   for row in rows:
      print(" "*indent + (" "*pad).join(line.ljust(col_width) for line in row))

if __name__ == '__main__':
    # STEP: Parse the command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", default=os.path.join(os.path.expanduser('~'), '.vdcapi'),
                    help="path/name of the config file to be used for the API URL and API keys (default is ~/.vdcapi)")
    parser.add_argument("-m", "--mode", choices=['deploy','d','print','p'], default='print',
                    help="specify the output mode: deploy (deploy the VM) or print (print the API call) (default: print)")
    parser.add_argument("-f", "--format", choices=['json','j','cloudmonkey','c','url'], default='json',
                        help="specify the format of the printed API call: json, cloudmonkey or url (default: json)")
    parser.add_argument("-r", "--region", choices=['Europe','europe','USA','usa','Asia','asia'],
                    default='Europe', help="specify the VDC region: Europe, USA or Asia (default Europe)")
    parser.add_argument("-i", "--iso", action='store_true', help="deploy from an ISO image")
    parser.add_argument("-k", "--keys", action='store_true', help="ask for choice of SSH keys")
    parser.add_argument("-u", "--userdata", action='store_true', help="ask for input of userdata by filename")
    parser.add_argument("-a", "--affinity", action='store_true', help="ask for selection of affinity group(s)")
    vdcRegion = parser.parse_args().region
    config_file = parser.parse_args().config
    mode = parser.parse_args().mode
    printFormat = parser.parse_args().format
    deployFromISO = parser.parse_args().iso
    askForSSHKeys = parser.parse_args().keys
    askForUserdata = parser.parse_args().userdata
    askForAffinityGroup = parser.parse_args().affinity

    # STEP: If config file is found, read its content,
    # else query user for the URL, API key, Secret key
    if os.path.isfile(config_file):
        with open(config_file) as fh:
            data = fh.read()
            config = json.loads(data)
            api_url = config['api_url']
            apiKey = config['api_key']
            secret = config['api_secret']
    else:
        print('API url (e.g. http://10.220.18.115:8080/client/api):', end='')
        api_url = raw_input()
        print('API key:', end='')
        apiKey = raw_input()
        secret = getpass.getpass(prompt='API secret:')

    # STEP: Create the API access object
    api = vdc.VDCApiCall(api_url, apiKey, secret)

#TO ADD: CHOOSE REGION OR FIND ZONES FOR ALL REGIONS THAT THE VDC ACCOUNT CAN ACCESS..........................

    # STEP: Select the zone
    result = api.listZones({'region': vdcRegion})
    zonelist = [zone['name'] for zone in result['zone']]
    zone_ids = [zone['id'] for zone in result['zone']]
    choice = choose_item_from_list(zonelist, prompt="Select the zone?")
    zone_id = zone_ids[choice['itemindex']]
    print("Selected zone: %s, %s\n" % (zone_id, choice['itemcontent']))

    # STEP: Select the template or ISO image
    if deployFromISO:
       print("Error: ISO case not implemented yet")
       exit
    else:
       request = {
        'region': vdcRegion,
        'templatefilter': 'executable',
        'zoneid': zone_id,
       }
       result = api.listTemplates(request)
       templates_sorted = sorted(result['template'], key=lambda item: item['name'].upper())
       templatelist = [template['name'] for template in templates_sorted]
       template_ids = [template['id'] for template in templates_sorted]
       choice = choose_item_from_list(templatelist, prompt="Select the template?")
       template_id = template_ids[choice['itemindex']]
       print("Selected template: %s, '%s'\n" % (template_id, choice['itemcontent']))

    # STEP: Select the compute/service offering
    ramlist = [512,1024,2048,4096,6144,8192,16384,24576,32768,65536,131072]
    choice_cpu = input("Input the number of CPUs? (between 1 and 12): ")
    print('')
    choice_ram = choose_item_from_list(map(lambda x: float(x)/1024, ramlist), prompt="Select the amount of RAM (GBytes)?")
    result = api.listServiceOfferings({'region': vdcRegion})
    serviceoffering_id = [s['id'] for s in result['serviceoffering'] if s['name']=='%d-%d'% (ramlist[choice_ram['itemindex']],choice_cpu)][0]
    print("Selected service offering: %s, \'%s\'\n" % (serviceoffering_id, '%d-%d'% (ramlist[choice_ram['itemindex']],choice_cpu)))

    # (optional) STEP: Select the affinity groups (if any exist)
    if askForAffinityGroup:
       print("Error: Select affinity groups not implemented yet")

    # STEP: Select the network(s)
    result = api.listNetworks({'region': vdcRegion, 'zoneid': zone_id})
    network_ids = [network['id'] for network in result['network']]
    networklist = ['%s (%s)' % (network['displaytext'],network['name']) for network in result['network']]
    network_num = input("Input the number of networks?: ")
    if network_num == 1:
       choice = choose_item_from_list(networklist, prompt="Select the network from the list? (this will be the default)")
       network_id = network_ids[choice['itemindex']]
       print("Selected network: %s, %s\n" % (network_id, choice['itemcontent']))
    else:
       choice = choose_item_from_list(networklist, prompt="Select network 1? (this will be the default)")
       network_id = [network_ids[choice['itemindex']]]
       for i in range(1,network_num):
          choice = choose_item_from_list(networklist, prompt="Select network %d?"%(i+1))
          network_id = network_id + [network_ids[choice['itemindex']]]
       print("Selected networks: %s\n" % (network_id))

    # (optional) STEP: Select keys
    if askForSSHKeys:
       result = api.listSSHKeyPairs({'region': vdcRegion})
       if result=={}:
          print("Error: SSH Keypair asked for, but none are set up in this account in this region")
          askForSSHKeys = False
       else:
          keypairnames = [keypair['name'] for keypair in result['sshkeypair']]
          keypairlist = ['%s (%s)' % (keypair['name'],keypair['fingerprint']) for keypair in result['sshkeypair']]
          choice = choose_item_from_list(keypairlist, prompt="Select SSH keypair?")
          keypairname = keypairnames[choice['itemindex']]

    # (optional) STEP: Input userdata
    if askForUserdata:
       print('')
       userdataFilename = raw_input('Input the path/name for the PLAIN TEXT file with the userdata: ')
       try:
          with open(userdataFilename) as fh:
             userdata=fh.read()
          userdata_b64 = base64.b64encode(userdata)
       except IOError:
          print("Error: Userdata plain text file not found or cannot be opened. Userdata will not be included in the deploy.")
          askForUserdata = False
          pass

    # STEP: Enter the VM 'name' and 'displaytext'
    ## TO DO: (1) create a default VM name, (2) do a test to check new name is unique against existing VMs
    #default_hostname =
    print('')
    hostname = raw_input(
       # 'Input name for the new VM (default %s):' % default_hostname
       'Input name for the new VM (name must be unique): '
    )
    hostname = hostname.strip()
    #if len(hostname) == 0:
    #    hostname = default_hostname

    default_displayname = hostname
    displayname = raw_input(
        'Input displayname for new VM (press ENTER for default: %s):' % default_displayname
    )
    displayname = displayname.strip()
    if len(displayname) == 0:
        displayname = default_displayname

    # STEP: Deploy the VM, or print out the API call
    deploy_params = {
        'region': vdcRegion,
        'serviceofferingid': serviceoffering_id,
        'templateid': template_id,
        'zoneid': zone_id,
        'displayname': displayname,
        'name': hostname,
        'networkids': network_id
    }
    if askForSSHKeys:
        deploy_params['keypair'] = keypairname
    if askForUserdata:
        deploy_params['userdata'] = userdata_b64
    print('')
    if mode=='print':
        print("-------------------------------------------\nRequired deploy command in %s format\n-------------------------------------------\n" % (printFormat))
        if printFormat=='json':
           print(json.dumps(deploy_params))
           print('')
        elif printFormat=='cloudmonkey':
           # for Cloudmonkey, 'region' is set by a separate command
           deploy_params.pop('region')
           params = ["%s=%s" % (key, deploy_params[key]) for key in deploy_params]
           print("NOTE: you need to execute the command 'set region %s' for the following deploy command to work...\n" % (vdcRegion))
           print("deploy virtualmachine " + " ".join(params))
           print('')
        else:
           deploy_params['command'] = 'deployVirtualMachine'
           deploy_params['response'] = 'json'
           deploy_params['apiKey'] = apiKey
           httprequest = zip(deploy_params.keys(), deploy_params.values())
           httprequest.sort(key=lambda x: x[0].lower())
           httprequest_data = "&".join(["=".join([r[0], urllib.quote_plus(str(r[1]))])
                                for r in httprequest])
           hashStr = "&".join(
               [
                "=".join(
                    [r[0].lower(),
                     str.lower(urllib.quote_plus(str(r[1]))).replace(
                         "+", "%20"
                     )]
                ) for r in httprequest
            ])
           sig = urllib.quote_plus(base64.b64encode(
             hmac.new(
                 secret,
                 hashStr,
                 hashlib.sha1
             ).digest()
           ).strip())
           httprequest_data += "&signature=%s" % sig
           print(api_url + "?" + httprequest_data)
           print('')
    else:
        print("Ready to deploy VM with these parameters:")
        pprint.pprint(json.dumps(deploy_params))
        choice = raw_input("Input D to deploy or any other key to exit:")
        if choice=="d" or choice=="D":
           result = api.deployVirtualMachine(deploy_params)
           job_id = result['jobid']
           pprint.pprint(api.wait_for_job(job_id))
<|endoftext|>


<commit_before>from httplib import HTTPConnection, HTTPMessage
from cStringIO import StringIO
from .files import save_cassette, load_cassette
from .cassette import Cassette

class VCRHTTPResponse(object):
    def __init__(self, recorded_response):
        self.recorded_response = recorded_response
        self.reason = recorded_response['status']['message']
        self.status = recorded_response['status']['code']
        self._content = StringIO(self.recorded_response['body']['string'])

        self.msg = HTTPMessage(StringIO(''))
        for k, v in self.recorded_response['headers'].iteritems():
            self.msg.addheader(k, v)

    def read(self, chunked=False):
        return self._content.read()

class VCRHTTPConnection(HTTPConnection):

    def __init__(self, *args, **kwargs):
        self._cassette = Cassette()
        HTTPConnection.__init__(self, *args, **kwargs)

    def _save_cassette(self):
        save_cassette(self._vcr_cassette_path, self._cassette)

    def request(self, method, url, body=None, headers={}):
        old_cassette = load_cassette(self._vcr_cassette_path)
        if old_cassette:
            return
        self._cassette.requests.append(dict(
            method=method,
            url=url,
            body=body,
            headers=headers
        ))
        return HTTPConnection.request(self, method, url, body=body, headers=headers)

    def getresponse(self, buffering=False):
        old_cassette = load_cassette(self._vcr_cassette_path)
        if old_cassette:
            return VCRHTTPResponse(old_cassette[0]['response'])
        response = HTTPConnection.getresponse(self)
        self._cassette.responses.append({
            'status': {'code': response.status, 'message': response.reason},
            'headers': dict(response.getheaders()),
            'body': {'string': response.read()},
        })
        self._save_cassette()
        return VCRHTTPResponse(self._cassette[0]['response'])

class VCRHTTPSConnection(VCRHTTPConnection):
    pass
<commit_msg>fix stupid bug?<commit_after>from httplib import HTTPConnection, HTTPMessage
from cStringIO import StringIO
from .files import save_cassette, load_cassette
from .cassette import Cassette

class VCRHTTPResponse(object):
    def __init__(self, recorded_response):
        self.recorded_response = recorded_response
        self.reason = recorded_response['status']['message']
        self.status = recorded_response['status']['code']
        self._content = StringIO(self.recorded_response['body']['string'])

        self.msg = HTTPMessage(StringIO(''))
        for k, v in self.recorded_response['headers'].iteritems():
            self.msg.addheader(k, v)

    def read(self, chunked=False):
        return self._content.read()

class VCRHTTPConnection(HTTPConnection):

    def __init__(self, *args, **kwargs):
        self._cassette = Cassette()
        HTTPConnection.__init__(self, *args, **kwargs)

    def _save_cassette(self):
        save_cassette(self._vcr_cassette_path, self._cassette)

    def request(self, method, url, body=None, headers={}):
        old_cassette = load_cassette(self._vcr_cassette_path)
        if old_cassette:
            return
        self._cassette.requests.append(dict(
            method=method,
            url=url,
            body=body,
            headers=headers
        ))
        return HTTPConnection.request(self, method, url, body=body, headers=headers)

    def getresponse(self, buffering=False):
        old_cassette = load_cassette(self._vcr_cassette_path)
        if not old_cassette:
            response = HTTPConnection.getresponse(self)
            self._cassette.responses.append({
                'status': {'code': response.status, 'message': response.reason},
                'headers': dict(response.getheaders()),
                'body': {'string': response.read()},
            })
            self._save_cassette()
        old_cassette = load_cassette(self._vcr_cassette_path)
        return VCRHTTPResponse(old_cassette[0]['response'])

class VCRHTTPSConnection(VCRHTTPConnection):
    pass
<|endoftext|>


<commit_before><commit_msg>Set DOWNLOAD_DELAY = 1, otherwise, stackoverflow.com will return 503 error when request too quickly.<commit_after><|endoftext|>


<commit_before>from __future__ import division
# Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from past.utils import old_div
from builtins import object
import logging
import posixpath
import time

from desktop.lib.i18n import smart_unicode
from desktop.lib.apputil import WARN_LEVEL_CALL_DURATION_MS, INFO_LEVEL_CALL_DURATION_MS

from desktop import conf

LOG = logging.getLogger(__name__)

class Resource(object):
  """
  Encapsulates a resource, and provides actions to invoke on it.
  """
  def __init__(self, client, relpath="", urlencode=True):
    """
    @param client: A Client object.
    @param relpath: The relative path of the resource.
    @param urlencode: percent encode paths.
    """
    self._client = client
    self._path = relpath.strip('/')
    self._urlencode = urlencode

  @property
  def base_url(self):
    return self._client.base_url

  def _join_uri(self, relpath):
    if relpath is None:
      return self._path
    return self._path + posixpath.normpath('/' + relpath)

  def _format_response(self, resp):
    """
    Decide whether the body should be a json dict or string
    """

    if resp.headers.get('location') and resp.headers.get('location').startswith('http://localhost:8080/'):
      return resp.headers.get('location').split('sessions/')[1]

    if len(resp.content) != 0 and resp.headers.get('content-type') and \
          'application/json' in resp.headers.get('content-type'):
      try:
        return resp.json()
      except Exception as ex:
        self._client.logger.exception('JSON decode error: %s' % resp.content)
        raise ex
    else:
      return resp.content

  def invoke(self, method, relpath=None, params=None, data=None, headers=None, files=None, allow_redirects=False, clear_cookies=False, log_response=True):
    resp = self._invoke(method=method,
                        relpath=relpath,
                        params=params,
                        data=data,
                        headers=headers,
                        files=files,
                        allow_redirects=allow_redirects,
                        clear_cookies=clear_cookies,
                        log_response=log_response)

    return self._format_response(resp)

  def _invoke(self, method, relpath=None, params=None, data=None, headers=None, files=None, allow_redirects=False, clear_cookies=False, log_response=True):
    """
    Invoke an API method.
    @return: Raw body or JSON dictionary (if response content type is JSON).
    """
    path = self._join_uri(relpath)
    start_time = time.time()
    resp = None
    try:
      resp = self._client.execute(method,
                                  path,
                                  params=params,
                                  data=data,
                                  headers=headers,
                                  files=files,
                                  allow_redirects=allow_redirects,
                                  urlencode=self._urlencode,
                                  clear_cookies=clear_cookies)
    finally: # Print the response time even when there's an exception
      log_length = conf.REST_RESPONSE_SIZE.get() != -1 and conf.REST_RESPONSE_SIZE.get() if log_response else 0 # We want to output duration without content
      duration = time.time() - start_time
      message = "%s %s Got response%s: %s%s" % (
          method,
          smart_unicode(path, errors='ignore'),
          ' in %dms' % (duration * 1000),
          smart_unicode(resp.content[:log_length], errors='replace') if resp else "",
          log_length and len(resp.content) > log_length and "..." or "" if resp else ""
      )
      self._client.logger.disabled = 0
      log_if_slow_call(duration=duration, message=message, logger=self._client.logger)

    return resp

  def get(self, relpath=None, params=None, headers=None, clear_cookies=False):
    """
    Invoke the GET method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("GET", relpath, params, headers=headers, allow_redirects=True, clear_cookies=clear_cookies)

  def delete(self, relpath=None, params=None, headers=None, clear_cookies=False):
    """
    Invoke the DELETE method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param headers: Optional. Base set of headers.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("DELETE", relpath, params, headers=headers, clear_cookies=clear_cookies)

  def post(self, relpath=None, params=None, data=None, contenttype=None, headers=None, files=None, allow_redirects=False,
           clear_cookies=False, log_response=True):
    """
    Invoke the POST method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param data: Optional. Body of the request.
    @param contenttype: Optional.
    @param headers: Optional. Base set of headers.
    @param allow_redirects: Optional. Allow request to automatically resolve redirects.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("POST", relpath, params, data, headers=self._make_headers(contenttype, headers), files=files,
                       allow_redirects=allow_redirects, clear_cookies=clear_cookies, log_response=log_response)

  def put(self, relpath=None, params=None, data=None, contenttype=None, allow_redirects=False, clear_cookies=False, headers=None):
    """
    Invoke the PUT method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param data: Optional. Body of the request.
    @param contenttype: Optional.
    @param allow_redirects: Optional. Allow request to automatically resolve redirects.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("PUT", relpath, params, data, headers=self._make_headers(contenttype, headers),
                       allow_redirects=allow_redirects, clear_cookies=clear_cookies)

  def _make_headers(self, contenttype=None, headers=None):
    if headers is None:
      headers = {}

    if contenttype:
      headers.update({'Content-Type': contenttype})

    return headers

  def resolve_redirect_url(self, method="GET", relpath=None, params=None, data=None, headers=None, files=None, allow_redirects=True, clear_cookies=False, log_response=True):
    resp = self._invoke(method=method,
                        relpath=relpath,
                        params=params,
                        data=data,
                        headers=headers,
                        files=files,
                        allow_redirects=allow_redirects,
                        clear_cookies=clear_cookies,
                        log_response=log_response)

    return resp.url.encode("utf-8")

# Same in thrift_util.py for not losing the trace class
def log_if_slow_call(duration, message, logger):
  if duration >= old_div(WARN_LEVEL_CALL_DURATION_MS, 1000):
    logger.warn('SLOW: %.2f - %s' % (duration, message))
  elif duration >= old_div(INFO_LEVEL_CALL_DURATION_MS, 1000):
    logger.info('SLOW: %.2f - %s' % (duration, message))
  else:
    logger.debug(message)
<commit_msg>HUE-8968 [core] Print more information on outbound HTTP request<commit_after>from __future__ import division
# Licensed to Cloudera, Inc. under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  Cloudera, Inc. licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from past.utils import old_div
from builtins import object
import logging
import posixpath
import urllib
import time

from django.utils.encoding import iri_to_uri, smart_str
from django.utils.http import urlencode

from desktop.lib.i18n import smart_unicode
from desktop.lib.apputil import WARN_LEVEL_CALL_DURATION_MS, INFO_LEVEL_CALL_DURATION_MS

from desktop import conf

LOG = logging.getLogger(__name__)

class Resource(object):
  """
  Encapsulates a resource, and provides actions to invoke on it.
  """
  def __init__(self, client, relpath="", urlencode=True):
    """
    @param client: A Client object.
    @param relpath: The relative path of the resource.
    @param urlencode: percent encode paths.
    """
    self._client = client
    self._path = relpath.strip('/')
    self._urlencode = urlencode

  @property
  def base_url(self):
    return self._client.base_url

  def _join_uri(self, relpath):
    if relpath is None:
      return self._path
    return self._path + posixpath.normpath('/' + relpath)

  def _format_response(self, resp):
    """
    Decide whether the body should be a json dict or string
    """

    if resp.headers.get('location') and resp.headers.get('location').startswith('http://localhost:8080/'):
      return resp.headers.get('location').split('sessions/')[1]

    if len(resp.content) != 0 and resp.headers.get('content-type') and \
          'application/json' in resp.headers.get('content-type'):
      try:
        return resp.json()
      except Exception as ex:
        self._client.logger.exception('JSON decode error: %s' % resp.content)
        raise ex
    else:
      return resp.content

  def invoke(self, method, relpath=None, params=None, data=None, headers=None, files=None, allow_redirects=False, clear_cookies=False, log_response=True):
    resp = self._invoke(method=method,
                        relpath=relpath,
                        params=params,
                        data=data,
                        headers=headers,
                        files=files,
                        allow_redirects=allow_redirects,
                        clear_cookies=clear_cookies,
                        log_response=log_response)

    return self._format_response(resp)

  def _invoke(self, method, relpath=None, params=None, data=None, headers=None, files=None, allow_redirects=False, clear_cookies=False, log_response=True):
    """
    Invoke an API method.
    @return: Raw body or JSON dictionary (if response content type is JSON).
    """
    path = self._join_uri(relpath)
    start_time = time.time()
    resp = None
    try:
      resp = self._client.execute(method,
                                  path,
                                  params=params,
                                  data=data,
                                  headers=headers,
                                  files=files,
                                  allow_redirects=allow_redirects,
                                  urlencode=self._urlencode,
                                  clear_cookies=clear_cookies)
    finally: # Print even when there's an exception
      log_length = conf.REST_RESPONSE_SIZE.get() != -1 and conf.REST_RESPONSE_SIZE.get() if log_response else 0 # We want to output duration without content
      duration = time.time() - start_time
      message = '%s %s %s%s%s %s%s returned in %dms %s %s %s%s' % (
        method,
        type(self._client._session.auth) if self._client._session and self._client._session.auth else None,
        self._client._base_url,
        smart_str(path),
        iri_to_uri('?' + urlencode(params)) if params else '',
        smart_unicode(data, errors='replace')[:log_length] if data else "",
        log_length and len(data) > log_length and "..." or "" if data else "",
        (duration * 1000),
        resp.status_code if resp else 0,
        len(resp.content) if resp else 0,
        smart_unicode(resp.content[:log_length], errors='replace') if resp else "",
        log_length and len(resp.content) > log_length and "..." or "" if resp else ""
      )
      self._client.logger.disabled = 0
      log_if_slow_call(duration=duration, message=message, logger=self._client.logger)

    return resp

  def get(self, relpath=None, params=None, headers=None, clear_cookies=False):
    """
    Invoke the GET method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("GET", relpath, params, headers=headers, allow_redirects=True, clear_cookies=clear_cookies)

  def delete(self, relpath=None, params=None, headers=None, clear_cookies=False):
    """
    Invoke the DELETE method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param headers: Optional. Base set of headers.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("DELETE", relpath, params, headers=headers, clear_cookies=clear_cookies)

  def post(self, relpath=None, params=None, data=None, contenttype=None, headers=None, files=None, allow_redirects=False,
           clear_cookies=False, log_response=True):
    """
    Invoke the POST method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param data: Optional. Body of the request.
    @param contenttype: Optional.
    @param headers: Optional. Base set of headers.
    @param allow_redirects: Optional. Allow request to automatically resolve redirects.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("POST", relpath, params, data, headers=self._make_headers(contenttype, headers), files=files,
                       allow_redirects=allow_redirects, clear_cookies=clear_cookies, log_response=log_response)

  def put(self, relpath=None, params=None, data=None, contenttype=None, allow_redirects=False, clear_cookies=False, headers=None):
    """
    Invoke the PUT method on a resource.
    @param relpath: Optional. A relative path to this resource's path.
    @param params: Key-value data.
    @param data: Optional. Body of the request.
    @param contenttype: Optional.
    @param allow_redirects: Optional. Allow request to automatically resolve redirects.
    @param clear_cookies: Optional. A flag to force any session cookies to be reset on the request.

    @return: A dictionary of the JSON result.
    """
    return self.invoke("PUT", relpath, params, data, headers=self._make_headers(contenttype, headers),
                       allow_redirects=allow_redirects, clear_cookies=clear_cookies)

  def _make_headers(self, contenttype=None, headers=None):
    if headers is None:
      headers = {}

    if contenttype:
      headers.update({'Content-Type': contenttype})

    return headers

  def resolve_redirect_url(self, method="GET", relpath=None, params=None, data=None, headers=None, files=None, allow_redirects=True, clear_cookies=False, log_response=True):
    resp = self._invoke(method=method,
                        relpath=relpath,
                        params=params,
                        data=data,
                        headers=headers,
                        files=files,
                        allow_redirects=allow_redirects,
                        clear_cookies=clear_cookies,
                        log_response=log_response)

    return resp.url.encode("utf-8")

# Same in thrift_util.py for not losing the trace class
def log_if_slow_call(duration, message, logger):
  if duration >= old_div(WARN_LEVEL_CALL_DURATION_MS, 1000):
    logger.warn('SLOW: %.2f - %s' % (duration, message))
  elif duration >= old_div(INFO_LEVEL_CALL_DURATION_MS, 1000):
    logger.info('SLOW: %.2f - %s' % (duration, message))
  else:
    logger.debug(message)
<|endoftext|>


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

import eventlet
import time
import socket
import fcntl
import struct
import threading

from oslo.config import cfg

from vmthunder.openstack.common import log as logging
from vmthunder.path import connection_to_str
from vmthunder.path import Path
from vmthunder.enum import Enum
from vmthunder.drivers import fcg
from vmthunder.drivers import dmsetup
from vmthunder.drivers import iscsi
from vmthunder.drivers import volt

CONF = cfg.CONF
LOG = logging.getLogger(__name__)

STATUS = Enum(['empty', 'building', 'ok', 'destroying', 'error'])
ACTIONS = Enum(['build', 'destroy'])

class Session(object):
    def __init__(self, volume_name):
        self.volume_name = volume_name
        self.root = {}
        self.paths = {}
        self.cached_path = ''
        self.has_multipath = False
        self.has_cache = False
        self.has_origin = False
        self.has_target = False
        self.is_login = False
        #TODO: all virtual machines called image
        self.vm = []
        self.peer_id = ''
        self.target_id = 0
        self.__status = STATUS.empty
        self.status_lock = threading.Lock()
        LOG.debug("VMThunder: create a session of volume_name %s" % self.volume_name)

    @property
    def origin_name(self):
        return 'origin_' + self.volume_name

    @property
    def origin_path(self):
        return dmsetup.prefix + self.origin_name

    @property
    def multipath_name(self):
        return 'multipath_' + self.volume_name

    @property
    def multipath_path(self):
        return dmsetup.prefix + self.multipath_name

    def change_status(self, src_status, dst_status):
        with self.status_lock:
            ret = False
            if self.__status == src_status:
                self.__status = dst_status
                ret = True
            LOG.debug("VMThunder: source status = %s, dst status = %s, ret = %s" % (src_status, dst_status, ret))
            return ret

    def deploy_image(self, image_connection):
        success = self.change_status(STATUS.empty, STATUS.building)
        if not success:
            while self.__status == STATUS.building:
                LOG.debug("VMThunder: in deploy_image, sleep 3 seconds waiting for build completed")
                eventlet.sleep(3)
        LOG.debug("VMThunder: ..........begin to deploy base image")
        try:
            origin_path = self._deploy_image(image_connection)
        except Exception, e:
            LOG.error(e)
            self.change_status(STATUS.building, STATUS.error)
            raise
        else:
            self.change_status(STATUS.building, STATUS.ok)
        LOG.debug("VMThunder: ..........deploy base image completed")
        return origin_path

    def _deploy_image(self, image_connection):
        #TODO: Roll back if failed !
        """
        deploy image in compute node, return the origin path to create snapshot
        :param image_connection: the connection towards to the base image
        :return: origin path to create snapshot
        """
        LOG.debug("VMThunder: in deploy_image, volume name = %s, has origin = %s, has cache = %s, "
                  "is_login = %s" % (self.volume_name, self.has_origin, self.has_cache, self.is_login))

        #Check current status
        if self.has_origin:
            return self.origin_path

        image_path = Path(self.reform_connection(image_connection))
        self.root[str(image_path)] = image_path

        LOG.debug("VMThunder: begin to rebuild paths")
        if len(self.paths) == 0:
            parent_list = self._get_parent()
            self.rebuild_paths(parent_list)
        LOG.debug("VMThunder: rebuild paths completed, multipath = %s" % self.multipath_path)
        if not self.has_cache:
            #TODO: NEED to fix here
            LOG.debug("VMThunder: create cache for base image %s" % self.volume_name)
            self.cached_path = self._create_cache(self.multipath_path)
            LOG.debug("VMThunder: create cache completed, cache path = %s" % self.cached_path)

        if not self.has_origin:
            LOG.debug("VMThunder: start to create origin, cache path = %s" % self.cached_path)
            self._create_origin(self.cached_path)
            LOG.debug("VMThunder: create origin complete, cache path = %s" % self.cached_path)

        LOG.debug("VMThunder: create target is an option!!!")

        if not self.has_target:
            LOG.debug("VMThunder: start to create target, cache path = %s" % self.cached_path)
            iqn = image_connection['target_iqn']
            self._create_target(iqn, self.cached_path)
            LOG.debug("VMThunder: create target complete, cache path = %s" % self.cached_path)

        return self.origin_path

    def destroy(self):
        LOG.debug("VMThunder: destroy session %s" % self.volume_name)
        assert not self.has_vm(), 'Destroy session %s failed, still has vm' % self.volume_name
        if self.is_login is True:
            volt.logout(self.volume_name, peer_id=self.peer_id)
            self.is_login = False
        if self.has_target:
            if iscsi.is_connected(self.target_id):
                return False
            else:
                self._delete_target()

        if self.has_origin:
            self._delete_origin()
        time.sleep(1)
        multipath = self.multipath_path
        if not self.has_origin and not self.has_target:
            self._delete_cache(multipath)
        if not self.has_cache:
            self._delete_multipath()
        if not self.has_multipath:
            for key in self.paths.keys():
                self.paths[key].disconnect()
                del self.paths[key]
        return True

    def adjust_for_heartbeat(self, parent_list):
        self.rebuild_paths(parent_list)
        LOG.debug('VMThunder: adjust_for_heartbeat according to connections:')
        LOG.debug(parent_list)

    def has_vm(self):
        if len(self.vm) > 0:
            return True
        else:
            return False

    def add_vm(self, vm_name):
        if vm_name not in self.vm:
            self.vm.append(vm_name)
        else:
            LOG.error("Add vm failed, VM %s existed" % vm_name)

    def rm_vm(self, vm_name):
        try:
            self.vm.remove(vm_name)
        except ValueError:
            LOG.error("remove vm failed. VM %s does not existed" % vm_name)

    def rebuild_paths(self, parents_list):
        #Reform connections
        parent_connections = self.reform_connections(parents_list)
        #Get keys of paths to remove
        keys_to_remove = []
        for key in self.paths.keys():
            found = False
            for connection in parent_connections:
                if str(self.paths[key]) == connection_to_str(connection):
                    found = True
                    break
            if not found:
                keys_to_remove.append(key)
        #If no path to connect, connect to root
        if len(parents_list) == 0:
            self.paths = self.root
        #Add paths to self.paths
        else:
            for parent in parent_connections:
                if isinstance(parent, dict):
                    parent_str = connection_to_str(parent)
                    if parent_str not in self.paths.keys():
                        self.paths[parent_str] = Path(parent)
                elif isinstance(parent, Path):
                    if str(parent) not in self.paths.keys():
                        self.paths[str(parent)] = parent
                else:
                    raise(Exception("Unknown %s type of %s "%(type(parent), parent)))
        #Connect new paths
        for key in self.paths.keys():
            if key not in keys_to_remove and not self.paths[key].connected:
                self.paths[key].connect()
        #Rebuild multipath device
        disks = [self.paths[key].device_path for key in self.paths
                 if key not in keys_to_remove and self.paths[key].connected]
        if self.has_multipath:
            self._reload_multipath(disks)
        else:
            self._create_multipath(disks)

        #Disconnect path to remove
        for key in keys_to_remove:
            if self.paths[key].connected:
                self.paths[key].disconnect()
            del self.paths[key]

    @staticmethod
    def _get_ip_address(ifname):
        LOG.debug("acquire ip address of %s" % ifname)
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        return socket.inet_ntoa(fcntl.ioctl(
            s.fileno(),
            0x8915,
            struct.pack('256s', ifname[:15]))[20:24])

    @staticmethod
    def reform_connection(connection):
        LOG.debug("old connection is :")
        LOG.debug(connection)
        new_connection = {'target_portal': connection['target_portal'],
                          'target_iqn': connection['target_iqn'],
                          'target_lun': connection['target_lun'],
        }
        LOG.debug("new connection is :")
        LOG.debug(new_connection)
        return new_connection

    def reform_connections(self, connections):
        new_connections = []
        for connection in connections:
            new_connections.append(self.reform_connection(connection))
        return new_connections

    def _create_target(self, iqn, path):
        if iscsi.exists(iqn) is False:
            self.target_id = iscsi.create_iscsi_target(iqn, path)
            LOG.debug("VMThunder: create a target and it's id is %s" % self.target_id)
            self.has_target = True
        #don't dynamic gain host_id and host_port
        #TODO: eth0? br100?
        host_ip = self._get_ip_address('br100')
        LOG.debug("VMThunder: try to login to master server")
        #TODO: port? lun? what is info
        if not self.is_login:
            info = volt.login(session_name=self.volume_name, peer_id=self.peer_id,
                              host=host_ip, port='3260', iqn=iqn, lun='1')
            LOG.debug("VMThunder: login to master server %s" % info)
            self.is_login = True

    def _delete_target(self):
        iscsi.remove_iscsi_target(0, 0, self.volume_name, self.volume_name)
        self.has_target = False
        LOG.debug("successful remove target %s " % self.target_id)

    def _create_multipath(self, disks):
        multipath_name = self.multipath_name
        multipath_path = dmsetup.multipath(multipath_name, disks)
        self.has_multipath = True
        LOG.debug("create multipath according connection :")
        LOG.debug(disks)
        return multipath_path

    def _delete_multipath(self):
        multipath_name = self.multipath_name
        dmsetup.remove_table(multipath_name)
        self.has_multipath = False
        LOG.debug("delete multipath of %s" % multipath_name)

    def _reload_multipath(self, disks):
        dmsetup.reload_multipath(self.multipath_name, disks)

    def _create_cache(self, multipath):
        cached_path = fcg.add_disk(multipath)
        self.has_cache = True
        LOG.debug("VMThunder: create cache according to multipath %s" % multipath)
        return cached_path

    def _delete_cache(self, multipath):
        fcg.rm_disk(multipath)
        self.has_cache = False
        LOG.debug("VMThunder: delete cache according to multipath %s " % multipath)

    def _create_origin(self, origin_dev):
        origin_name = self.origin_name
        if self.has_origin:
            origin_path = self.origin_path
        else:
            origin_path = dmsetup.origin(origin_name, origin_dev)
            LOG.debug("VMThunder: create origin on %s" % origin_dev)
            self.has_origin = True
        return origin_path

    def _delete_origin(self):
        origin_name = self.origin_name
        dmsetup.remove_table(origin_name)
        LOG.debug("VMThunder: remove origin %s " % origin_name)
        self.has_origin = False

    def _get_parent(self):
        #TODO: br100???
        host_ip = self._get_ip_address('br100')
        while True:
            self.peer_id, parent_list = volt.get(session_name=self.volume_name, host=host_ip)
            LOG.debug("VMThunder: in get_parent function to get parent_list :")
            LOG.debug(parent_list)
            #Wait for parents are ready
            bo = True
            for parent in parent_list:
                if parent.status == "pending":
                    bo = False
                    break
            if bo:
                return parent_list
            #TODO: sleep???
            time.sleep(1)
<commit_msg>change get_ip to conf<commit_after>#!/usr/bin/env python

import eventlet
import time
import socket
import fcntl
import struct
import threading

from oslo.config import cfg

from vmthunder.openstack.common import log as logging
from vmthunder.path import connection_to_str
from vmthunder.path import Path
from vmthunder.enum import Enum
from vmthunder.drivers import fcg
from vmthunder.drivers import dmsetup
from vmthunder.drivers import iscsi
from vmthunder.drivers import volt

CONF = cfg.CONF
LOG = logging.getLogger(__name__)

STATUS = Enum(['empty', 'building', 'ok', 'destroying', 'error'])
ACTIONS = Enum(['build', 'destroy'])

class Session(object):
    def __init__(self, volume_name):
        self.volume_name = volume_name
        self.root = {}
        self.paths = {}
        self.cached_path = ''
        self.has_multipath = False
        self.has_cache = False
        self.has_origin = False
        self.has_target = False
        self.is_login = False
        #TODO: all virtual machines called image
        self.vm = []
        self.peer_id = ''
        self.target_id = 0
        self.__status = STATUS.empty
        self.status_lock = threading.Lock()
        LOG.debug("VMThunder: create a session of volume_name %s" % self.volume_name)

    @property
    def origin_name(self):
        return 'origin_' + self.volume_name

    @property
    def origin_path(self):
        return dmsetup.prefix + self.origin_name

    @property
    def multipath_name(self):
        return 'multipath_' + self.volume_name

    @property
    def multipath_path(self):
        return dmsetup.prefix + self.multipath_name

    def change_status(self, src_status, dst_status):
        with self.status_lock:
            ret = False
            if self.__status == src_status:
                self.__status = dst_status
                ret = True
            LOG.debug("VMThunder: source status = %s, dst status = %s, ret = %s" % (src_status, dst_status, ret))
            return ret

    def deploy_image(self, image_connection):
        success = self.change_status(STATUS.empty, STATUS.building)
        if not success:
            while self.__status == STATUS.building:
                LOG.debug("VMThunder: in deploy_image, sleep 3 seconds waiting for build completed")
                eventlet.sleep(3)
        LOG.debug("VMThunder: ..........begin to deploy base image")
        try:
            origin_path = self._deploy_image(image_connection)
        except Exception, e:
            LOG.error(e)
            self.change_status(STATUS.building, STATUS.error)
            raise
        else:
            self.change_status(STATUS.building, STATUS.ok)
        LOG.debug("VMThunder: ..........deploy base image completed")
        return origin_path

    def _deploy_image(self, image_connection):
        #TODO: Roll back if failed !
        """
        deploy image in compute node, return the origin path to create snapshot
        :param image_connection: the connection towards to the base image
        :return: origin path to create snapshot
        """
        LOG.debug("VMThunder: in deploy_image, volume name = %s, has origin = %s, has cache = %s, "
                  "is_login = %s" % (self.volume_name, self.has_origin, self.has_cache, self.is_login))

        #Check current status
        if self.has_origin:
            return self.origin_path

        image_path = Path(self.reform_connection(image_connection))
        self.root[str(image_path)] = image_path

        LOG.debug("VMThunder: begin to rebuild paths")
        if len(self.paths) == 0:
            parent_list = self._get_parent()
            self.rebuild_paths(parent_list)
        LOG.debug("VMThunder: rebuild paths completed, multipath = %s" % self.multipath_path)
        if not self.has_cache:
            #TODO: NEED to fix here
            LOG.debug("VMThunder: create cache for base image %s" % self.volume_name)
            self.cached_path = self._create_cache(self.multipath_path)
            LOG.debug("VMThunder: create cache completed, cache path = %s" % self.cached_path)

        if not self.has_origin:
            LOG.debug("VMThunder: start to create origin, cache path = %s" % self.cached_path)
            self._create_origin(self.cached_path)
            LOG.debug("VMThunder: create origin complete, cache path = %s" % self.cached_path)

        LOG.debug("VMThunder: create target is an option!!!")

        if not self.has_target:
            LOG.debug("VMThunder: start to create target, cache path = %s" % self.cached_path)
            iqn = image_connection['target_iqn']
            self._create_target(iqn, self.cached_path)
            LOG.debug("VMThunder: create target complete, cache path = %s" % self.cached_path)

        return self.origin_path

    def destroy(self):
        LOG.debug("VMThunder: destroy session %s" % self.volume_name)
        assert not self.has_vm(), 'Destroy session %s failed, still has vm' % self.volume_name
        if self.is_login is True:
            volt.logout(self.volume_name, peer_id=self.peer_id)
            self.is_login = False
        if self.has_target:
            if iscsi.is_connected(self.target_id):
                return False
            else:
                self._delete_target()

        if self.has_origin:
            self._delete_origin()
        time.sleep(1)
        multipath = self.multipath_path
        if not self.has_origin and not self.has_target:
            self._delete_cache(multipath)
        if not self.has_cache:
            self._delete_multipath()
        if not self.has_multipath:
            for key in self.paths.keys():
                self.paths[key].disconnect()
                del self.paths[key]
        return True

    def adjust_for_heartbeat(self, parent_list):
        self.rebuild_paths(parent_list)
        LOG.debug('VMThunder: adjust_for_heartbeat according to connections:')
        LOG.debug(parent_list)

    def has_vm(self):
        if len(self.vm) > 0:
            return True
        else:
            return False

    def add_vm(self, vm_name):
        if vm_name not in self.vm:
            self.vm.append(vm_name)
        else:
            LOG.error("Add vm failed, VM %s existed" % vm_name)

    def rm_vm(self, vm_name):
        try:
            self.vm.remove(vm_name)
        except ValueError:
            LOG.error("remove vm failed. VM %s does not existed" % vm_name)

    def rebuild_paths(self, parents_list):
        #Reform connections
        parent_connections = self.reform_connections(parents_list)
        #Get keys of paths to remove
        keys_to_remove = []
        for key in self.paths.keys():
            found = False
            for connection in parent_connections:
                if str(self.paths[key]) == connection_to_str(connection):
                    found = True
                    break
            if not found:
                keys_to_remove.append(key)
        #If no path to connect, connect to root
        if len(parents_list) == 0:
            self.paths = self.root
        #Add paths to self.paths
        else:
            for parent in parent_connections:
                if isinstance(parent, dict):
                    parent_str = connection_to_str(parent)
                    if parent_str not in self.paths.keys():
                        self.paths[parent_str] = Path(parent)
                elif isinstance(parent, Path):
                    if str(parent) not in self.paths.keys():
                        self.paths[str(parent)] = parent
                else:
                    raise(Exception("Unknown %s type of %s "%(type(parent), parent)))
        #Connect new paths
        for key in self.paths.keys():
            if key not in keys_to_remove and not self.paths[key].connected:
                self.paths[key].connect()
        #Rebuild multipath device
        disks = [self.paths[key].device_path for key in self.paths
                 if key not in keys_to_remove and self.paths[key].connected]
        if self.has_multipath:
            self._reload_multipath(disks)
        else:
            self._create_multipath(disks)

        #Disconnect path to remove
        for key in keys_to_remove:
            if self.paths[key].connected:
                self.paths[key].disconnect()
            del self.paths[key]

    @staticmethod
    def _get_ip_address(ifname):
        LOG.debug("acquire ip address of %s" % ifname)
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        return socket.inet_ntoa(fcntl.ioctl(
            s.fileno(),
            0x8915,
            struct.pack('256s', ifname[:15]))[20:24])

    @staticmethod
    def reform_connection(connection):
        LOG.debug("old connection is :")
        LOG.debug(connection)
        new_connection = {'target_portal': connection['target_portal'],
                          'target_iqn': connection['target_iqn'],
                          'target_lun': connection['target_lun'],
        }
        LOG.debug("new connection is :")
        LOG.debug(new_connection)
        return new_connection

    def reform_connections(self, connections):
        new_connections = []
        for connection in connections:
            new_connections.append(self.reform_connection(connection))
        return new_connections

    def _create_target(self, iqn, path):
        if iscsi.exists(iqn) is False:
            self.target_id = iscsi.create_iscsi_target(iqn, path)
            LOG.debug("VMThunder: create a target and it's id is %s" % self.target_id)
            self.has_target = True
        #don't dynamic gain host_id and host_port
        #TODO: eth0? br100?
        host_ip = self._get_ip_address('br100')
        LOG.debug("VMThunder: try to login to master server")
        #TODO: port? lun? what is info
        if not self.is_login:
            info = volt.login(session_name=self.volume_name, peer_id=self.peer_id,
                              host=host_ip, port='3260', iqn=iqn, lun='1')
            LOG.debug("VMThunder: login to master server %s" % info)
            self.is_login = True

    def _delete_target(self):
        iscsi.remove_iscsi_target(0, 0, self.volume_name, self.volume_name)
        self.has_target = False
        LOG.debug("successful remove target %s " % self.target_id)

    def _create_multipath(self, disks):
        multipath_name = self.multipath_name
        multipath_path = dmsetup.multipath(multipath_name, disks)
        self.has_multipath = True
        LOG.debug("create multipath according connection :")
        LOG.debug(disks)
        return multipath_path

    def _delete_multipath(self):
        multipath_name = self.multipath_name
        dmsetup.remove_table(multipath_name)
        self.has_multipath = False
        LOG.debug("delete multipath of %s" % multipath_name)

    def _reload_multipath(self, disks):
        dmsetup.reload_multipath(self.multipath_name, disks)

    def _create_cache(self, multipath):
        cached_path = fcg.add_disk(multipath)
        self.has_cache = True
        LOG.debug("VMThunder: create cache according to multipath %s" % multipath)
        return cached_path

    def _delete_cache(self, multipath):
        fcg.rm_disk(multipath)
        self.has_cache = False
        LOG.debug("VMThunder: delete cache according to multipath %s " % multipath)

    def _create_origin(self, origin_dev):
        origin_name = self.origin_name
        if self.has_origin:
            origin_path = self.origin_path
        else:
            origin_path = dmsetup.origin(origin_name, origin_dev)
            LOG.debug("VMThunder: create origin on %s" % origin_dev)
            self.has_origin = True
        return origin_path

    def _delete_origin(self):
        origin_name = self.origin_name
        dmsetup.remove_table(origin_name)
        LOG.debug("VMThunder: remove origin %s " % origin_name)
        self.has_origin = False

    def _get_parent(self):
        #TODO: br100???
        #host_ip = self._get_ip_address('br100')
        host_ip = CONF.host_ip
        while True:
            self.peer_id, parent_list = volt.get(session_name=self.volume_name, host=host_ip)
            LOG.debug("VMThunder: in get_parent function to get parent_list :")
            LOG.debug(parent_list)
            #Wait for parents are ready
            bo = True
            for parent in parent_list:
                if parent.status == "pending":
                    bo = False
                    break
            if bo:
                return parent_list
            #TODO: sleep???
            time.sleep(1)
<|endoftext|>


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

'''
Twitch.py - Twitch on the CLI
See readme.md for more info
'''

__author__  = 'Steven Leibrock'
__email__   = 'leibrockoli@gmail.com'
__version__ = '1.1.4'

from argparse import ArgumentParser
from subprocess import call
from unittest import main as TestMain
from unittest import TestCase

# Constants
API_URL                = 'https://api.twitch.tv/kraken'
TOP_GAMES              = '/games/top'
SEARCH_STR             = '/search/streams'
DEFAULT_LIMIT          = 10
LIVESTREAMER_INSTALLED = True
REQUESTS_INSTALLED     = True
CHAR_LIMIT             = 50
DEBUG                  = False

# Livestream qualities
# -q will disable the automatic use of AUTO_QUALITY
USE_AUTO_QUALITY       = True
AUTO_QUALITY           = 'source'

# If Livestreamer is installed, will check to figure out where it is
# Livestreamer installs are in different places based on the OS
LIVESTREAMER_PATH = '/usr/bin/livestreamer'
POSSIBLE_PATHS = ['/usr/bin/livestreamer', '/usr/local/bin/livestreamer']

# Try loading the Requests library
try:
    if debug: print("Checking for Requests...")
    from requests import get as re_get
except ImportError as e:
    print('Error: {0}'.format(e))
    print('Requests not installed, to install it you can run:')
    print('\n  pip install requests\n')
    print('Defaulting to urllib...')
    from urllib import urlopen
    from json import loads as json_load
    REQUESTS_INSTALLED = False

# Try loading Livestreamer (if failed, tell them how to get)
try:
    if debug: print("Checking for Livestreamer...")
    from livestreamer import streams as StreamData
    from livestreamer import __version__ as LSV
    if debug: print('Livestreamer version: {v}'.format(v=LSV))

    # Try to locate where Livestreamer is installed
    from os.path import isfile as where_is_file
    for path in POSSIBLE_PATHS:
        if where_is_file(path):
            LIVESTREAMER_PATH = path
except Exception as e:
    print('Error: {0}'.format(e))
    print('Livestreamer not installed, to install it you can run:')
    print('\n  pip install livestreamer\n')
    LIVESTREAMER_INSTALLED = False

def limit_string(s):
    '''
    Limit a string to the CHAR_LIMIT
    :s is the input string
    '''
    # Daft Punk song variable for extra credit
    too_long = ".." if len(s) > CHAR_LIMIT else ""
    return "{0}{1}".format(s[:CHAR_LIMIT], too_long)

def load_url(url):
    '''
    Return a dictionary of the JSON
    This function will either use urllib or requests
    based on what's installed
    '''
    if REQUESTS_INSTALLED:
        return re_get(url).json()
    return json_load(urlopen(url).read().decode('utf-8'))

def get_input(disp_str, err_str, maxv):
    '''
    Fetch input from the user and return the number given
    Subtract 1 at the end to make it from [0:(maxv-1)]
    '''
    inp = None
    while inp is None:
        try:
            inp = int(input(disp_str))
            if inp < 1 or inp > maxv:
                inp = None
        except ValueError:
            print(err_str)
    return inp - 1

def load_into_livestreamer(url):
    '''
    Load a URL into Livestreamer
    :url is the target URL to load
    '''
    print('Loading {0} ...\n'.format(url))

    # Check for USE_AUTO_QUALITY
    if USE_AUTO_QUALITY:
        print('Using auto-quality "{0}" ...\n'.format(AUTO_QUALITY))
        q = AUTO_QUALITY
    else:
        # Scan the URL for qualities
        qualities = StreamData(url)
        names = [q.lower() for q in qualities.keys()]
        for i, q in enumerate(names):
            print('[{ind}] {qual}'.format(ind=i+1, qual=q))
        inp = get_input('Enter quality> ', 'Try again', len(names))
        print('Using quality "{0}" ...\n'.format(names[inp]))
        q = names[inp]
    call([LIVESTREAMER_PATH, url, q])
    return True

def main_directory(limit):
    '''
    Fetch the top games on Twitch
    :limit determines how many results will be fetched
    '''
    print('Fetching main Twitch.tv directory...')
    print('Press CTRL+C to quit\n')
    url = '{0}{1}?limit={2}'.format(API_URL, TOP_GAMES, limit)
    blob = load_url(url)
    longest = max([len(g['game']['name']) for g in blob['top']])
    game_titles = [g['game']['name'] for g in blob['top']]
    for i, game_blob in enumerate(blob['top']):
        print('{n:>{fill2}}) {gt:<{fill}}   [viewers: {vc:7,}, chans: {cc:5,}]'.
              format(gt=game_blob['game']['name'], vc=game_blob['viewers'],
                     cc=game_blob['channels'], fill=longest, n=i+1,
                     fill2=len(str(DEFAULT_LIMIT))))
    inp = get_input('Enter number> ', 'Try again', len(game_titles))
    return scan_game_directory(game_titles[inp], limit)

def scan_game_directory(game, limit):
    '''
    Scan game directory
    :game is the game streams to find
    :limit is how many to fetch
    This code also includes input to select Stream to open
    '''
    print('Getting streams for "{g}"'.format(g=game))
    print('Press CTRL+C to quit\n')
    search_key = game.lower().replace(' ', '%20')
    search_url = '{b}{q}?q=%22{g}%22&limit={l}'.format(
        b=API_URL, q=SEARCH_STR, g=search_key, l=limit)
    blob = load_url(search_url)
    urls = [u['channel']['url'] for u in blob['streams']]
    highest_views = max([len("{0:,}".format(v['viewers'])) for v in
                         blob['streams']])
    if DEBUG:
        print('Total streams: {0}'.format(blob['_total']))
        print('Total fetched: {0}\n'.format(len(blob['streams'])))
    for i, stream in enumerate(blob['streams']):
        print('{n:>{f}}) {s: <{cl}}  [viewers: {v:{mv},}]'.format(
                s=limit_string(stream['channel']['status']),
                n=i+1, v=stream['viewers'], f=len(str(DEFAULT_LIMIT)),
                cl=CHAR_LIMIT+2, mv=highest_views))

    if not LIVESTREAMER_INSTALLED:
        print('\n{0}\n!! Livestreamer is not installed !!\n{0}'.format('*'*10))
        return False

    inp = get_input('Select stream> ', 'Try again', len(urls))
    return load_into_livestreamer(urls[inp])

class TwitchTest(TestCase):

    def test_front_page(self):
        pass

    def test_limit_count(self):
        pass

if __name__ == '__main__':
    parser = ArgumentParser(
        description='Interact with Twitch.tv',
        prog='twitch.py')
    parser.add_argument('-g', type=str, nargs='+', metavar='text',
                        help='The game whose directory you wish to scan')
    parser.add_argument('-l', '--limit', type=int, default=DEFAULT_LIMIT,
                        metavar='LIM', help='Number of streams to fetch')
    parser.add_argument('-d', '--debug', help='Debug the program',
                        action='store_const', default=False, const=True)
    parser.add_argument('-v', '--version',action='version',
                        version='%(prog)s ver.{0}'.format(__version__))
    parser.add_argument('-q', const=True, default=False,
                        action='store_const', help='Disable auto-quality')
    parser.add_argument('-t', '--test', action='store_const', const=True,
                        default=False, help="Run a unit test then quit")
    try:
        args = parser.parse_args()
        DEBUG = args.debug
        USE_AUTO_QUALITY = args.q
        if DEBUG:
            print(args)

        if args.test:
            print("Running Unit Test")
            quit()

        if args.game is None or args.game == '':
            main_directory(args.limit)
        else:
            scan_game_directory(' '.join(args.game), args.limit)
    except KeyboardInterrupt as e:
        print('\nQuitting...')
    except Exception as e:
        print("Error encountered: {0}".format(e))
    print('\nAll done!')

# end
<commit_msg>Fixed debug not working and invalid args<commit_after>#!/usr/bin/python3
# -*- coding: utf-8 -*-

'''
Twitch.py - Twitch on the CLI
See readme.md for more info
'''

__author__  = 'Steven Leibrock'
__email__   = 'leibrockoli@gmail.com'
__version__ = '1.1.5'

from argparse import ArgumentParser
from subprocess import call
from unittest import main as TestMain
from unittest import TestCase

# Constants
API_URL                = 'https://api.twitch.tv/kraken'
TOP_GAMES              = '/games/top'
SEARCH_STR             = '/search/streams'
DEFAULT_LIMIT          = 10
LIVESTREAMER_INSTALLED = True
REQUESTS_INSTALLED     = True
CHAR_LIMIT             = 50
DEBUG                  = False

# Livestream qualities
# -q will disable the automatic use of AUTO_QUALITY
USE_AUTO_QUALITY       = True
AUTO_QUALITY           = 'source'

# If Livestreamer is installed, will check to figure out where it is
# Livestreamer installs are in different places based on the OS
LIVESTREAMER_PATH = '/usr/bin/livestreamer'
POSSIBLE_PATHS = ['/usr/bin/livestreamer', '/usr/local/bin/livestreamer']

# Try loading the Requests library
try:
    if DEBUG: print("Checking for Requests...")
    from requests import get as re_get
except ImportError as e:
    print('Error: {0}'.format(e))
    print('Requests not installed, to install it you can run:')
    print('\n  pip install requests\n')
    print('Defaulting to urllib...')
    from urllib import urlopen
    from json import loads as json_load
    REQUESTS_INSTALLED = False

# Try loading Livestreamer (if failed, tell them how to get)
try:
    if DEBUG: print("Checking for Livestreamer...")
    from livestreamer import streams as StreamData
    from livestreamer import __version__ as LSV
    if DEBUG: print('Livestreamer version: {v}'.format(v=LSV))

    # Try to locate where Livestreamer is installed
    from os.path import isfile as where_is_file
    for path in POSSIBLE_PATHS:
        if where_is_file(path):
            LIVESTREAMER_PATH = path
except Exception as e:
    print('Error: {0}'.format(e))
    print('Livestreamer not installed, to install it you can run:')
    print('\n  pip install livestreamer\n')
    LIVESTREAMER_INSTALLED = False

def limit_string(s):
    '''
    Limit a string to the CHAR_LIMIT
    :s is the input string
    '''
    # Daft Punk song variable for extra credit
    too_long = ".." if len(s) > CHAR_LIMIT else ""
    return "{0}{1}".format(s[:CHAR_LIMIT], too_long)

def load_url(url):
    '''
    Return a dictionary of the JSON
    This function will either use urllib or requests
    based on what's installed
    '''
    if REQUESTS_INSTALLED:
        return re_get(url).json()
    return json_load(urlopen(url).read().decode('utf-8'))

def get_input(disp_str, err_str, maxv):
    '''
    Fetch input from the user and return the number given
    Subtract 1 at the end to make it from [0:(maxv-1)]
    '''
    inp = None
    while inp is None:
        try:
            inp = int(input(disp_str))
            if inp < 1 or inp > maxv:
                inp = None
        except ValueError:
            print(err_str)
    return inp - 1

def load_into_livestreamer(url):
    '''
    Load a URL into Livestreamer
    :url is the target URL to load
    '''
    print('Loading {0} ...\n'.format(url))

    # Check for USE_AUTO_QUALITY
    if USE_AUTO_QUALITY:
        print('Using auto-quality "{0}" ...\n'.format(AUTO_QUALITY))
        q = AUTO_QUALITY
    else:
        # Scan the URL for qualities
        qualities = StreamData(url)
        names = [q.lower() for q in qualities.keys()]
        for i, q in enumerate(names):
            print('[{ind}] {qual}'.format(ind=i+1, qual=q))
        inp = get_input('Enter quality> ', 'Try again', len(names))
        print('Using quality "{0}" ...\n'.format(names[inp]))
        q = names[inp]
    call([LIVESTREAMER_PATH, url, q])
    return True

def main_directory(limit):
    '''
    Fetch the top games on Twitch
    :limit determines how many results will be fetched
    '''
    print('Fetching main Twitch.tv directory...')
    print('Press CTRL+C to quit\n')
    url = '{0}{1}?limit={2}'.format(API_URL, TOP_GAMES, limit)
    blob = load_url(url)
    longest = max([len(g['game']['name']) for g in blob['top']])
    game_titles = [g['game']['name'] for g in blob['top']]
    for i, game_blob in enumerate(blob['top']):
        print('{n:>{fill2}}) {gt:<{fill}}   [viewers: {vc:7,}, chans: {cc:5,}]'.
              format(gt=game_blob['game']['name'], vc=game_blob['viewers'],
                     cc=game_blob['channels'], fill=longest, n=i+1,
                     fill2=len(str(DEFAULT_LIMIT))))
    inp = get_input('Enter number> ', 'Try again', len(game_titles))
    return scan_game_directory(game_titles[inp], limit)

def scan_game_directory(game, limit):
    '''
    Scan game directory
    :game is the game streams to find
    :limit is how many to fetch
    This code also includes input to select Stream to open
    '''
    print('Getting streams for "{g}"'.format(g=game))
    print('Press CTRL+C to quit\n')
    search_key = game.lower().replace(' ', '%20')
    search_url = '{b}{q}?q=%22{g}%22&limit={l}'.format(
        b=API_URL, q=SEARCH_STR, g=search_key, l=limit)
    blob = load_url(search_url)
    urls = [u['channel']['url'] for u in blob['streams']]
    highest_views = max([len("{0:,}".format(v['viewers'])) for v in
                         blob['streams']])
    if DEBUG:
        print('Total streams: {0}'.format(blob['_total']))
        print('Total fetched: {0}\n'.format(len(blob['streams'])))
    for i, stream in enumerate(blob['streams']):
        print('{n:>{f}}) {s: <{cl}}  [viewers: {v:{mv},}]'.format(
                s=limit_string(stream['channel']['status']),
                n=i+1, v=stream['viewers'], f=len(str(DEFAULT_LIMIT)),
                cl=CHAR_LIMIT+2, mv=highest_views))

    if not LIVESTREAMER_INSTALLED:
        print('\n{0}\n!! Livestreamer is not installed !!\n{0}'.format('*'*10))
        return False

    inp = get_input('Select stream> ', 'Try again', len(urls))
    return load_into_livestreamer(urls[inp])

class TwitchTest(TestCase):

    def test_front_page(self):
        pass

    def test_limit_count(self):
        pass

if __name__ == '__main__':
    parser = ArgumentParser(
        description='Interact with Twitch.tv',
        prog='twitch.py')
    parser.add_argument('-g', type=str, nargs='+', metavar='text',
                        help='The game whose directory you wish to scan')
    parser.add_argument('-l', '--limit', type=int, default=DEFAULT_LIMIT,
                        metavar='LIM', help='Number of streams to fetch')
    parser.add_argument('-d', '--debug', help='Debug the program',
                        action='store_const', default=False, const=True)
    parser.add_argument('-v', '--version',action='version',
                        version='%(prog)s ver.{0}'.format(__version__))
    parser.add_argument('-q', const=True, default=False,
                        action='store_const', help='Disable auto-quality')
    parser.add_argument('-t', '--test', action='store_const', const=True,
                        default=False, help="Run a unit test then quit")
    try:
        args = parser.parse_args()
        DEBUG = args.debug
        USE_AUTO_QUALITY = args.q
        if DEBUG:
            print(args)

        if args.test:
            print("Running Unit Test")
            quit()

        if args.g is None or args.g == '':
            main_directory(args.limit)
        else:
            scan_game_directory(' '.join(args.g), args.limit)
    except KeyboardInterrupt as e:
        print('\nQuitting...')
    except Exception as e:
        print("Error encountered: {0}".format(e))
    print('\nAll done!')

# end
<|endoftext|>