<commit_before>#!/usr/bin/python
#
# DVB Transport stream should pick out entire DVB services and send those to
# named outboxes. (Send DVB Services to a Kamaelia Service...(!))
#

import os
import dvb3.frontend
import dvb3.dmx
import time
import struct

from Axon.Component import component
DVB_PACKET_SIZE = 188
DVB_RESYNC = "\x47"
import Axon.AdaptiveCommsComponent

def tune_DVBT(fe, frequency):
    # Build the tuning parameters (DVB-T)
    params = dvb3.frontend.OFDMParameters()
    params.frequency = frequency * 1000 * 1000

    params.inversion = dvb3.frontend.INVERSION_AUTO
    params.constellation = dvb3.frontend.QAM_16
    params.coderate_HP = dvb3.frontend.FEC_3_4
    params.coderate_LP = dvb3.frontend.FEC_3_4

    # Start the tuning
    fe.set_frontend(params)

def notLocked(fe):
    """\
    Wait for lock, if it's not available, yield a true value.
    If it is, exit with a StopIteration. (allows use in a for
    loop)
    """
    return (fe.read_status() & dvb3.frontend.FE_HAS_LOCK) != 0

def addPIDS(pids):
    """\
    Adds the given PID to the transport stream that will be available
    in "/dev/dvb/adapter0/dvr0"
    """
    demuxers = [dvb3.dmx.Demux(0, blocking = 0) for _ in pids]
    for p in xrange(len(pids)):
        demuxers[p].set_pes_filter(pids[p],
                                   dvb3.dmx.DMX_IN_FRONTEND,
                                   dvb3.dmx.DMX_OUT_TS_TAP,
                                   dvb3.dmx.DMX_PES_OTHER,
                                   dvb3.dmx.DMX_IMMEDIATE_START)
    return demuxers

class DVB_Multiplex(component):
    """\
    This is a DVB Multiplex Tuner.

    This tunes the given DVB card to the given frequency. This then sets
    up the dvr0 device node to recieve the data recieved on a number of
    PIDs.

    A special case use of these is to tune to 2 specific PIDs - the audio
    and video for a specific TV channel. If you pass just 2 PIDs then
    you're tuning to a specific channel.

    NOTE 1: This multiplex tuner deliberately does not know what
    frequency the multiplex is on, and does not know what PIDs are
    inside that multiplex. You are expected to find out this information
    independently.

    NOTE 2: This means also that producing a mock for the next stages in
    this system should be relatively simple - we run this code once and
    dump to disk.
    """
    def __init__(self, freq, pids):
        self.freq = freq
        self.pids = pids
        super(DVB_Multiplex, self).__init__()

    def main(self):
        # Open the frontend of card 0 (/dev/dvb/adaptor0/frontend0)
        fe = dvb3.frontend.Frontend(0, blocking=0)
        tune_DVBT(fe, self.freq)
        while notLocked(fe): yield 1  # could sleep for, say, 0.1 seconds.
        demuxers = addPIDS(self.pids)

        # This is then a file reader, actually.
        # Should be a little more system friendly really
        fd = os.open("/dev/dvb/adapter0/dvr0", os.O_RDONLY | os.O_NONBLOCK)
        while True:
            try:
               data = os.read(fd, 2048)
               self.send(data, "outbox")
            except OSError:
               pass
            yield 1
            # XXX: We should add the following:
            # XXX: Handle shutdown messages
            # XXX: Pass on shutdown messages/errors

class DVB_Demuxer(Axon.AdaptiveCommsComponent.AdaptiveCommsComponent):
    """\
    This demuxer expects to recieve the output from a DVB_Multiplex
    component on its primary inbox. It is also provided with a number
    of pids. For each pid that it knows about, it forwards the data
    received on that PID to an appropriate outbox. Data associated with
    unknown PIDs in the datastream is thrown away.

    The output here is still transport stream packets. Another layer
    is required to decide what to do with these - to yank out the PES
    and ES packets.
    """
    Inboxes = {
        "inbox" : "This is where we expect to recieve a transport stream",
        "control" : "We will receive shutdown messages here",
    }
    def __init__(self, pidmap):
        super(DVB_Demuxer, self).__init__()
        self.pidmap = pidmap
        for pid in pidmap: # This adds an outbox per pid
            # This allows for the PIDs to be split or remultiplexed
            # together.
            if not self.outboxes.has_key(pidmap[pid]):
                self.addOutbox(pidmap[pid])

    def errorIndicatorSet(self, packet):  return ord(packet[1]) & 0x80
    def scrambledPacket(self, packet):    return ord(packet[3]) & 0xc0

    def main(self):
        buffer = ""
        while 1:
            yield 1
            if self.dataReady("inbox"):
              buffer += self.recv("inbox")

            while len(buffer) >= DVB_PACKET_SIZE:
                  yield 1
                  i = buffer.find(DVB_RESYNC)
                  if i == -1: # if not found
                      "we have a dud"
                      buffer = ""
                      continue
                  if i>0:
                      # if found remove all bytes preceeding that point in the buffers
                      # And try again
                      buffer = buffer[i:]
                      continue
                  # packet is the first 188 bytes in the buffer now
                  packet, buffer = buffer[:DVB_PACKET_SIZE], buffer[DVB_PACKET_SIZE:]

                  if self.errorIndicatorSet(packet): continue
                  if self.scrambledPacket(packet):   continue

                  pid = struct.unpack(">H", packet[1: 3])[0] & 0x1fff

                  # Send the packet to the outbox appropriate for this PID.
                  # "Fail" silently for PIDs we don't know about and weren't
                  # asked to demultiplex
                  try:
#                      print ".", self.pidmap, str(pid)
                      self.send(packet, self.pidmap[ str(pid) ])
                  except KeyError:
                      pass

#
# XXX
#
# This is where we may wish to think about piping the results to something like
# mencoder, perhaps via a named pipe, since that will *probably* work. We'll have to
# check that though :-/
#
# For XTech, we can see if we can coax Mencoder into taking our output direct
# and doing something useful/interesting... (At least then we're up and running)
#

class PIDHandler:

    counter = None
    pes_dirty = False
    pes = None

    def process_ts_packet(self, packet):
        x = ord(packet[3])
        counter = x & 0x0f
        adaptation = (x & 0x30) >> 4

        if adaptation == 0:
            # "ITU-T Rec. H2.222.0 | ISO/IEC 13818-1 decoders shall discard
            # Transport Stream packets with the adaptation_field_control field
            # set to a value of '00'".
            return

        if adaptation != 1 and ord(packet[4]) > 0:
            discontinuity = ((ord(packet[5]) & 0x80) == 0x80)
        else:
            discontinuity = False
        if self.counter is None:
            self.counter = counter
        elif adaptation & 1:
            self.counter = (self.counter + 1) & 0xf

        if self.counter != counter:
            if not discontinuity:
                self.pes_dirty = True
            self.counter = counter

        if ord(packet[1]) & 0x40:
            # Payload Unit Start Indicator (PUSI) is set, so a new packet has
            # started.  Throw the old one down to the PES processor.
            if self.pes is not None and not self.pes_dirty:
                self.process_pes_packet("".join(self.pes))
            self.pes_dirty = False
            self.pes = None

        if adaptation == 1:
            chunk = packet[4: ]
        elif adaptation == 3:
            chunk = packet[5 + ord(packet[4]): ]
        elif adaptation == 2:
            # No interesting data.
            return

        if len(chunk) > 0:
            if self.pes is None:
                self.pes = [chunk]
            else:
                self.pes.append(chunk)

if __name__ == "__main__":
    from Kamaelia.Util.PipelineComponent import pipeline
    from Kamaelia.File.Writing import SimpleFileWriter
    from Kamaelia.ReadFileAdaptor import ReadFileAdaptor
    from Kamaelia.Util.Graphline import Graphline

    channels_london =  {
           "MORE4+1" : (   538, #MHz
                         [ 701, 702 ] # PID (programme ID) for video and PID for audio
                       )
    }
    services = {
           "NEWS24": (754, [640, 641]),
           "MORE4+1": (810, [701,702]),
           "TMF": (810, [201,202])
    }
    if 0:
        pipeline(
           DVB_Multiplex(754, [640, 641, 620, 621, 622, 610, 611, 612, 600, 601, 602, 12]),
           SimpleFileWriter("multiplex_new.data")
        ).run()
    if 1:
        Graphline(
            SOURCE=ReadFileAdaptor("multiplex.data"),
            DEMUX=DVB_Demuxer({
                "640": "NEWS24",
                "641": "NEWS24",
                "600": "BBCONE",
                "601": "BBCONE",
                "610": "BBCTWO",
                "611": "BBCTWO",
                "620": "CBBC",
                "621": "CBBC",
            }),
            NEWS24=SimpleFileWriter("news24.data"),
            BBCONE=SimpleFileWriter("bbcone.data"),
            BBCTWO=SimpleFileWriter("bbctwo.data"),
            CBBC=SimpleFileWriter("cbbc.data"),
            linkages={
               ("SOURCE", "outbox"):("DEMUX","inbox"),
               ("DEMUX", "NEWS24"): ("NEWS24", "inbox"),
               ("DEMUX", "BBCONE"): ("BBCONE", "inbox"),
               ("DEMUX", "BBCTWO"): ("BBCTWO", "inbox"),
               ("DEMUX", "CBBC"): ("CBBC", "inbox"),
            }
        ).run()
#            if not self.outboxes.has_key(pidmap[pid]):
#                self.addOutbox(pidmap[pid])
<commit_msg>Changed DVB_Multiplex into a threadedcomponent that now blocks whilst waiting for data from the demuxer.<commit_after>#!/usr/bin/python
#
# DVB Transport stream should pick out entire DVB services and send those to
# named outboxes. (Send DVB Services to a Kamaelia Service...(!))
#

import os
import dvb3.frontend
import dvb3.dmx
import time
import struct

from Axon.Component import component
DVB_PACKET_SIZE = 188
DVB_RESYNC = "\x47"
import Axon.AdaptiveCommsComponent

def tune_DVBT(fe, frequency):
    # Build the tuning parameters (DVB-T)
    params = dvb3.frontend.OFDMParameters()
    params.frequency = frequency * 1000 * 1000

    params.inversion = dvb3.frontend.INVERSION_AUTO
    params.constellation = dvb3.frontend.QAM_16
    params.coderate_HP = dvb3.frontend.FEC_2_3
    params.coderate_LP = dvb3.frontend.FEC_2_3

    # Start the tuning
    fe.set_frontend(params)

def notLocked(fe):
    """\
    Wait for lock, if it's not available, yield a true value.
    If it is, exit with a StopIteration. (allows use in a for
    loop)
    """
    return (fe.read_status() & dvb3.frontend.FE_HAS_LOCK) != 0

def addPIDS(pids):
    """\
    Adds the given PID to the transport stream that will be available
    in "/dev/dvb/adapter0/dvr0"
    """
    demuxers = [dvb3.dmx.Demux(0, blocking = 0) for _ in pids]
    for p in xrange(len(pids)):
        demuxers[p].set_pes_filter(pids[p],
                                   dvb3.dmx.DMX_IN_FRONTEND,
                                   dvb3.dmx.DMX_OUT_TS_TAP,
                                   dvb3.dmx.DMX_PES_OTHER,
                                   dvb3.dmx.DMX_IMMEDIATE_START)
    return demuxers

from Axon.ThreadedComponent import threadedcomponent

class DVB_Multiplex(threadedcomponent):
    """\
    This is a DVB Multiplex Tuner.

    This tunes the given DVB card to the given frequency. This then sets
    up the dvr0 device node to recieve the data recieved on a number of
    PIDs.

    A special case use of these is to tune to 2 specific PIDs - the audio
    and video for a specific TV channel. If you pass just 2 PIDs then
    you're tuning to a specific channel.

    NOTE 1: This multiplex tuner deliberately does not know what
    frequency the multiplex is on, and does not know what PIDs are
    inside that multiplex. You are expected to find out this information
    independently.

    NOTE 2: This means also that producing a mock for the next stages in
    this system should be relatively simple - we run this code once and
    dump to disk.
    """
    def __init__(self, freq, pids):
        self.freq = freq
        self.pids = pids
        super(DVB_Multiplex, self).__init__()

    def main(self):
        # Open the frontend of card 0 (/dev/dvb/adaptor0/frontend0)
        fe = dvb3.frontend.Frontend(0, blocking=0)
        tune_DVBT(fe, self.freq)
        while notLocked(fe): time.sleep(0.1)  #yield 1  # could sleep for, say, 0.1 seconds.
        demuxers = addPIDS(self.pids)

        # This is then a file reader, actually.
        # Should be a little more system friendly really
        fd = os.open("/dev/dvb/adapter0/dvr0", os.O_RDONLY) # | os.O_NONBLOCK)
        while True:
            try:
               data = os.read(fd, 2048)
               self.send(data, "outbox")
            except OSError:
               pass
#            yield 1
            # XXX: We should add the following:
            # XXX: Handle shutdown messages
            # XXX: Pass on shutdown messages/errors

class DVB_Demuxer(Axon.AdaptiveCommsComponent.AdaptiveCommsComponent):
    """\
    This demuxer expects to recieve the output from a DVB_Multiplex
    component on its primary inbox. It is also provided with a number
    of pids. For each pid that it knows about, it forwards the data
    received on that PID to an appropriate outbox. Data associated with
    unknown PIDs in the datastream is thrown away.

    The output here is still transport stream packets. Another layer
    is required to decide what to do with these - to yank out the PES
    and ES packets.
    """
    Inboxes = {
        "inbox" : "This is where we expect to recieve a transport stream",
        "control" : "We will receive shutdown messages here",
    }
    def __init__(self, pidmap):
        super(DVB_Demuxer, self).__init__()
        self.pidmap = pidmap
        for pid in pidmap: # This adds an outbox per pid
            # This allows for the PIDs to be split or remultiplexed
            # together.
            if not self.outboxes.has_key(pidmap[pid]):
                self.addOutbox(pidmap[pid])

    def errorIndicatorSet(self, packet):  return ord(packet[1]) & 0x80
    def scrambledPacket(self, packet):    return ord(packet[3]) & 0xc0

    def main(self):
        buffer = ""
        while 1:
            yield 1
            while self.dataReady("inbox"):
              buffer += self.recv("inbox")

            while len(buffer) >= DVB_PACKET_SIZE:
                  yield 1
                  i = buffer.find(DVB_RESYNC)
                  if i == -1: # if not found
                      "we have a dud"
                      buffer = ""
                      continue
                  if i>0:
                      # if found remove all bytes preceeding that point in the buffers
                      # And try again
                      buffer = buffer[i:]
                      continue
                  # packet is the first 188 bytes in the buffer now
                  packet, buffer = buffer[:DVB_PACKET_SIZE], buffer[DVB_PACKET_SIZE:]

                  if self.errorIndicatorSet(packet): continue
                  if self.scrambledPacket(packet):   continue

                  pid = struct.unpack(">H", packet[1: 3])[0] & 0x1fff

                  # Send the packet to the outbox appropriate for this PID.
                  # "Fail" silently for PIDs we don't know about and weren't
                  # asked to demultiplex
                  try:
#                      print ".", self.pidmap, str(pid)
                      self.send(packet, self.pidmap[ str(pid) ])
                  except KeyError:
                      pass

            self.pause()

#
# XXX
#
# This is where we may wish to think about piping the results to something like
# mencoder, perhaps via a named pipe, since that will *probably* work. We'll have to
# check that though :-/
#
# For XTech, we can see if we can coax Mencoder into taking our output direct
# and doing something useful/interesting... (At least then we're up and running)
#

class PIDHandler:

    counter = None
    pes_dirty = False
    pes = None

    def process_ts_packet(self, packet):
        x = ord(packet[3])
        counter = x & 0x0f
        adaptation = (x & 0x30) >> 4

        if adaptation == 0:
            # "ITU-T Rec. H2.222.0 | ISO/IEC 13818-1 decoders shall discard
            # Transport Stream packets with the adaptation_field_control field
            # set to a value of '00'".
            return

        if adaptation != 1 and ord(packet[4]) > 0:
            discontinuity = ((ord(packet[5]) & 0x80) == 0x80)
        else:
            discontinuity = False
        if self.counter is None:
            self.counter = counter
        elif adaptation & 1:
            self.counter = (self.counter + 1) & 0xf

        if self.counter != counter:
            if not discontinuity:
                self.pes_dirty = True
            self.counter = counter

        if ord(packet[1]) & 0x40:
            # Payload Unit Start Indicator (PUSI) is set, so a new packet has
            # started.  Throw the old one down to the PES processor.
            if self.pes is not None and not self.pes_dirty:
                self.process_pes_packet("".join(self.pes))
            self.pes_dirty = False
            self.pes = None

        if adaptation == 1:
            chunk = packet[4: ]
        elif adaptation == 3:
            chunk = packet[5 + ord(packet[4]): ]
        elif adaptation == 2:
            # No interesting data.
            return

        if len(chunk) > 0:
            if self.pes is None:
                self.pes = [chunk]
            else:
                self.pes.append(chunk)

if __name__ == "__main__":
    from Kamaelia.Util.PipelineComponent import pipeline
    from Kamaelia.File.Writing import SimpleFileWriter
    from Kamaelia.ReadFileAdaptor import ReadFileAdaptor
    from Kamaelia.Util.Graphline import Graphline

    channels_london =  {
           "MORE4+1" : (   538, #MHz
                         [ 701, 702 ] # PID (programme ID) for video and PID for audio
                       )
    }
    services = {
           "NEWS24": (754, [640, 641]),
           "MORE4+1": (810, [701,702]),
           "TMF": (810, [201,202])
    }
    if 0:
        pipeline(
           DVB_Multiplex(754, [640, 641, 620, 621, 622, 610, 611, 612, 600, 601, 602, 12]),
           SimpleFileWriter("multiplex_new.data")
        ).run()
    if 1:
        Graphline(
            SOURCE=ReadFileAdaptor("multiplex.data"),
            DEMUX=DVB_Demuxer({
                "640": "NEWS24",
                "641": "NEWS24",
                "600": "BBCONE",
                "601": "BBCONE",
                "610": "BBCTWO",
                "611": "BBCTWO",
                "620": "CBBC",
                "621": "CBBC",
            }),
            NEWS24=SimpleFileWriter("news24.data"),
            BBCONE=SimpleFileWriter("bbcone.data"),
            BBCTWO=SimpleFileWriter("bbctwo.data"),
            CBBC=SimpleFileWriter("cbbc.data"),
            linkages={
               ("SOURCE", "outbox"):("DEMUX","inbox"),
               ("DEMUX", "NEWS24"): ("NEWS24", "inbox"),
               ("DEMUX", "BBCONE"): ("BBCONE", "inbox"),
               ("DEMUX", "BBCTWO"): ("BBCTWO", "inbox"),
               ("DEMUX", "CBBC"): ("CBBC", "inbox"),
            }
        ).run()
#            if not self.outboxes.has_key(pidmap[pid]):
#                self.addOutbox(pidmap[pid])
<|endoftext|>


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

import abc
import contextlib
import logging
import ntpath
import os
import time

import bs4
import six
from six.moves import urllib

from argus import exceptions
from argus import util

CONF = util.get_config()
LOG = util.get_logger()

__all__ = (
    'InstancePreparer',
    'WindowsInstancePreparer',
)

def _read_url(url):
    request = urllib.request.urlopen(url)
    with contextlib.closing(request) as stream:
        content = stream.read()
        if six.PY3:
            content = content.decode(errors='replace')
        return content

def _get_git_link():
    content = _read_url("http://git-scm.com/download/win")
    soup = bs4.BeautifulSoup(content)
    download_div = soup.find('div', {'class': 'callout downloading'})
    if not download_div:
        raise exceptions.CloudbaseCIError(
            "Could not find callout_downloading div.")

    for a_object in download_div.find_all('a'):
        href = a_object.get('href', '')
        if not href.endswith('.exe'):
            continue
        return href
    raise exceptions.CloudbaseCIError("git download link not found.")

@six.add_metaclass(abc.ABCMeta)
class InstancePreparer(object):
    """Handle instance preparing.

    The method :meth:`~prepare` does all the necessary work for
    preparing a new instance. The executed steps are:

    * wait for boot completion.
    * get an install script for CloudbaseInit
    * installs CloudbaseInit
    * waits for the finalization of the installation.
    """

    def __init__(self, instance_id, servers_client, remote_client):
        self._servers_client = servers_client
        self._instance_id = instance_id
        self._remote_client = remote_client

    def _execute(self, cmd):
        """Execute the given command and fail when the command fails."""
        stdout, stderr, return_code = self._remote_client.run_remote_cmd(cmd)
        if return_code:
            raise exceptions.CloudbaseCIError(
                "Command {command!r} failed with "
                "return code {return_code!r}"
                .format(command=cmd,
                        return_code=return_code))
        return stdout, stderr

    def _run_cmd_until_condition(self, cmd, cond, retry_count=None,
                                 retry_count_interval=5):
        """Run the given `cmd` until a condition *cond* occurs.

        :param cmd:
            A string, representing a command which needs to
            be executed on the underlying remote client.
        :param cond:
            A callable which receives the stdout returned by
            executing the command. It should return a boolean value,
            which tells to this function to stop execution.
        :param retry_count:
            The number of retries which this function has.
            If the value is ``None``, then the function will run *forever*.
        :param retry_count_interval:
            The number of seconds to sleep when retrying a command.
        """
        count = 0
        while True:
            try:
                std_out, std_err = self._execute(cmd)
            except Exception:
                LOG.debug("Command %r failed while waiting for condition",
                          cmd)
                count += 1
                if retry_count and count >= retry_count:
                    raise exceptions.CloudbaseTimeoutError(
                        "Command {!r} failed too many times."
                        .format(cmd))
                time.sleep(retry_count_interval)
            else:
                if std_err:
                    raise exceptions.CloudbaseCLIError(
                        "Executing command {!r} failed with {!r}"
                        .format(cmd, std_err))
                elif cond(std_out):
                    break
                else:
                    time.sleep(retry_count_interval)

    @abc.abstractmethod
    def wait_for_boot_completion(self):
        """Wait for the instance to finish up booting."""

    @abc.abstractmethod
    def get_installation_script(self):
        """Get the installation script for cloudbaseinit."""

    @abc.abstractmethod
    def install_cbinit(self):
        """Install the cloudbaseinit code."""

    @abc.abstractmethod
    def wait_cbinit_finalization(self):
        """Wait for the finalization of cloudbaseinit."""

    @abc.abstractmethod
    def wait_reboot(self):
        """Do a reboot and wait for the instance to be up."""

    @abc.abstractmethod
    def install_git(self):
        """Install git in the instance."""

    @abc.abstractmethod
    def sysprep(self):
        """Do the final steps after installing cloudbaseinit.

        This requires running sysprep on Windows, but on other
        platforms there might be no need for calling it.
        """

    @abc.abstractmethod
    def replace_code(self):
        """Do whatever is necessary to replace the code for cloudbaseinit."""

    def prepare(self):
        """Prepare the underlying instance.

        The following operations will be executed:

        * wait for boot completion
        * get an installation script for CloudbaseInit
        * install CloudbaseInit by running the previously downloaded file.
        * wait until the instance is up and running.
        """
        LOG.info("Preparing instance %s", self._instance_id)
        self.wait_for_boot_completion()
        self.get_installation_script()
        self.install_cbinit()
        self.install_git()

        if CONF.argus.replace_code:
            self.replace_code()

        self.sysprep()
        self.wait_reboot()
        self.wait_cbinit_finalization()
        LOG.info("Finished preparing instance %s", self._instance_id)

    if CONF.argus.debug:
        prepare = util.trap_failure(prepare)

class WindowsInstancePreparer(InstancePreparer):
    """Instance preparer for Windows machines."""

    def get_program_files(self):
        """Get the location of program files from the instance."""
        stdout, _ = self._execute('powershell "(Get-WmiObject  Win32_OperatingSystem).'
                                  'OSArchitecture"')
        architecture = stdout.strip()

        # Next, get the location.
        if architecture == '64-bit':
            location, _ = self._execute('powershell "${ENV:ProgramFiles(x86)}"')
        else:
            location, _ = self._execute('powershell "$ENV:ProgramFiles"')
        return location.strip()

    def wait_for_boot_completion(self):
        LOG.info("Waiting for boot completion")

        wait_cmd = ('powershell "(Get-WmiObject Win32_Account | '
                    'where -Property Name -contains {0}).Name"'
                    .format(CONF.argus.default_ci_username))
        return self._run_cmd_until_condition(
            wait_cmd,
            lambda stdout: stdout.strip() == CONF.argus.default_ci_username)

    def get_installation_script(self):
        """Get an insallation script for CloudbaseInit."""
        LOG.info("Retrieve an installation script for CloudbaseInit")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/installCBinit.ps1 -outfile C:\\installcbinit.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

    def install_cbinit(self):
        """Run the installation script for CloudbaseInit."""
        LOG.info("Run the downloaded installation script")

        cmd = ('powershell "C:\\\\installcbinit.ps1 -serviceType {}"'
               .format(CONF.argus.service_type))
        self._execute(cmd)

    def install_git(self):
        """Install git in the instance."""
        LOG.info("Installing git.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/install_git.ps1 -outfile C:\\\\install_git.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

        git_link = _get_git_link()
        git_base = os.path.basename(git_link)
        cmd = ('powershell "C:\\install_git.ps1 {} {}"'
               .format(git_link, git_base))
        self._execute(cmd)

    def replace_code(self):
        """Replace the code of cloudbaseinit."""
        LOG.info("Replacing cloudbaseinit's code.")

        LOG.info("Getting program files location.")
        # Get the program files location.
        program_files = self.get_program_files()

        # Remove everything from the cloudbaseinit installation.
        LOG.info("Removing recursively cloudbaseinit.")
        cloudbaseinit = ntpath.join(
            program_files, "Cloudbase Solutions",
            "Cloudbase-Init",
            # TODO(cpopa): take care of this when testing Python 3.
            "Python27",
            "Lib",
            "site-packages",
            "cloudbaseinit")
        self._execute('rmdir "{}" /S /q'.format(cloudbaseinit))

        # Clone the repo
        LOG.info("cloning the cloudbaseinit repo.")
        self._execute("git clone https://github.com/stackforge/"
                      "cloudbase-init C:\\cloudbaseinit")

        # Run the command provided at cli.
        LOG.info("Applying cli patch.")
        opts = util.parse_cli()

        if not opts.git_command:
            raise exceptions.CloudbaseCLIError(
                "git_command not passed at command line. "
                "Could not replace anything.")

        self._execute("cd C:\\cloudbaseinit && {}".format(opts.git_command))

        # Replace the code, by moving the code from cloudbaseinit
        # to the installed location.
        LOG.info("Replacing code.")
        self._execute('powershell "Copy-Item C:\\cloudbaseinit\\cloudbaseinit '
                      '\'{}\' -Recurse"'.format(cloudbaseinit))

    def sysprep(self):
        """Prepare the instance for the actual tests, by running sysprep."""
        LOG.info("Running sysprep.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/sysprep.ps1 -outfile 'C:\\sysprep.ps1'"
               .format(CONF.argus.resources))
        self._execute(cmd)
        self._execute('powershell C:\\sysprep.ps1')

    def wait_cbinit_finalization(self):
        """Wait for the finalization of CloudbaseInit.

        The function waits until all the plugins have been executed.
        """
        LOG.info("Waiting for the finalization of CloudbaseInit execution")

        # Test that this instance's cloudbaseinit run exists.
        key = ('HKLM:SOFTWARE\\Wow6432Node\\Cloudbase` '
               'Solutions\\Cloudbase-init\\{0}'
               .format(self._instance_id))
        self._run_cmd_until_condition(
            'powershell "Test-Path \'{0}\'"'.format(key),
            lambda out: out.strip() == 'True')

        # Test the number of executed cloudbaseinit plugins.
        wait_cmd = 'powershell (Get-Item %s\\Plugins).ValueCount' % key
        self._run_cmd_until_condition(
            wait_cmd,
            lambda out: int(out) >= int(CONF.argus.expected_plugins_count))

    def wait_reboot(self):
        """Do a reboot and wait until the instance is up."""

        LOG.info('Waiting for server status SHUTOFF because of sysprep')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='SHUTOFF',
            extra_timeout=600)

        self._servers_client.start(self._instance_id)

        LOG.info('Waiting for server status ACTIVE')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='ACTIVE')

<commit_msg>Fix the ridiculous quote wrap.<commit_after># Copyright 2014 Cloudbase-init
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
"""Instance preparing utilities."""

import abc
import contextlib
import logging
import ntpath
import os
import time

import bs4
import six
from six.moves import urllib

from argus import exceptions
from argus import util

CONF = util.get_config()
LOG = util.get_logger()

__all__ = (
    'InstancePreparer',
    'WindowsInstancePreparer',
)

def _read_url(url):
    request = urllib.request.urlopen(url)
    with contextlib.closing(request) as stream:
        content = stream.read()
        if six.PY3:
            content = content.decode(errors='replace')
        return content

def _get_git_link():
    content = _read_url("http://git-scm.com/download/win")
    soup = bs4.BeautifulSoup(content)
    download_div = soup.find('div', {'class': 'callout downloading'})
    if not download_div:
        raise exceptions.CloudbaseCIError(
            "Could not find callout_downloading div.")

    for a_object in download_div.find_all('a'):
        href = a_object.get('href', '')
        if not href.endswith('.exe'):
            continue
        return href
    raise exceptions.CloudbaseCIError("git download link not found.")

@six.add_metaclass(abc.ABCMeta)
class InstancePreparer(object):
    """Handle instance preparing.

    The method :meth:`~prepare` does all the necessary work for
    preparing a new instance. The executed steps are:

    * wait for boot completion.
    * get an install script for CloudbaseInit
    * installs CloudbaseInit
    * waits for the finalization of the installation.
    """

    def __init__(self, instance_id, servers_client, remote_client):
        self._servers_client = servers_client
        self._instance_id = instance_id
        self._remote_client = remote_client

    def _execute(self, cmd):
        """Execute the given command and fail when the command fails."""
        stdout, stderr, return_code = self._remote_client.run_remote_cmd(cmd)
        if return_code:
            raise exceptions.CloudbaseCIError(
                "Command {command!r} failed with "
                "return code {return_code!r}"
                .format(command=cmd,
                        return_code=return_code))
        return stdout, stderr

    def _run_cmd_until_condition(self, cmd, cond, retry_count=None,
                                 retry_count_interval=5):
        """Run the given `cmd` until a condition *cond* occurs.

        :param cmd:
            A string, representing a command which needs to
            be executed on the underlying remote client.
        :param cond:
            A callable which receives the stdout returned by
            executing the command. It should return a boolean value,
            which tells to this function to stop execution.
        :param retry_count:
            The number of retries which this function has.
            If the value is ``None``, then the function will run *forever*.
        :param retry_count_interval:
            The number of seconds to sleep when retrying a command.
        """
        count = 0
        while True:
            try:
                std_out, std_err = self._execute(cmd)
            except Exception:
                LOG.debug("Command %r failed while waiting for condition",
                          cmd)
                count += 1
                if retry_count and count >= retry_count:
                    raise exceptions.CloudbaseTimeoutError(
                        "Command {!r} failed too many times."
                        .format(cmd))
                time.sleep(retry_count_interval)
            else:
                if std_err:
                    raise exceptions.CloudbaseCLIError(
                        "Executing command {!r} failed with {!r}"
                        .format(cmd, std_err))
                elif cond(std_out):
                    break
                else:
                    time.sleep(retry_count_interval)

    @abc.abstractmethod
    def wait_for_boot_completion(self):
        """Wait for the instance to finish up booting."""

    @abc.abstractmethod
    def get_installation_script(self):
        """Get the installation script for cloudbaseinit."""

    @abc.abstractmethod
    def install_cbinit(self):
        """Install the cloudbaseinit code."""

    @abc.abstractmethod
    def wait_cbinit_finalization(self):
        """Wait for the finalization of cloudbaseinit."""

    @abc.abstractmethod
    def wait_reboot(self):
        """Do a reboot and wait for the instance to be up."""

    @abc.abstractmethod
    def install_git(self):
        """Install git in the instance."""

    @abc.abstractmethod
    def sysprep(self):
        """Do the final steps after installing cloudbaseinit.

        This requires running sysprep on Windows, but on other
        platforms there might be no need for calling it.
        """

    @abc.abstractmethod
    def replace_code(self):
        """Do whatever is necessary to replace the code for cloudbaseinit."""

    def prepare(self):
        """Prepare the underlying instance.

        The following operations will be executed:

        * wait for boot completion
        * get an installation script for CloudbaseInit
        * install CloudbaseInit by running the previously downloaded file.
        * wait until the instance is up and running.
        """
        LOG.info("Preparing instance %s", self._instance_id)
        self.wait_for_boot_completion()
        self.get_installation_script()
        self.install_cbinit()
        self.install_git()

        if CONF.argus.replace_code:
            self.replace_code()

        self.sysprep()
        self.wait_reboot()
        self.wait_cbinit_finalization()
        LOG.info("Finished preparing instance %s", self._instance_id)

    if CONF.argus.debug:
        prepare = util.trap_failure(prepare)

class WindowsInstancePreparer(InstancePreparer):
    """Instance preparer for Windows machines."""

    def get_program_files(self):
        """Get the location of program files from the instance."""
        stdout, _ = self._execute('powershell "(Get-WmiObject  Win32_OperatingSystem).'
                                  'OSArchitecture"')
        architecture = stdout.strip()

        # Next, get the location.
        if architecture == '64-bit':
            location, _ = self._execute('powershell "${ENV:ProgramFiles(x86)}"')
        else:
            location, _ = self._execute('powershell "$ENV:ProgramFiles"')
        return location.strip()

    def wait_for_boot_completion(self):
        LOG.info("Waiting for boot completion")

        wait_cmd = ('powershell "(Get-WmiObject Win32_Account | '
                    'where -Property Name -contains {0}).Name"'
                    .format(CONF.argus.default_ci_username))
        return self._run_cmd_until_condition(
            wait_cmd,
            lambda stdout: stdout.strip() == CONF.argus.default_ci_username)

    def get_installation_script(self):
        """Get an insallation script for CloudbaseInit."""
        LOG.info("Retrieve an installation script for CloudbaseInit")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/installCBinit.ps1 -outfile C:\\installcbinit.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

    def install_cbinit(self):
        """Run the installation script for CloudbaseInit."""
        LOG.info("Run the downloaded installation script")

        cmd = ('powershell "C:\\\\installcbinit.ps1 -serviceType {}"'
               .format(CONF.argus.service_type))
        self._execute(cmd)

    def install_git(self):
        """Install git in the instance."""
        LOG.info("Installing git.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/install_git.ps1 -outfile C:\\\\install_git.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

        git_link = _get_git_link()
        git_base = os.path.basename(git_link)
        cmd = ('powershell "C:\\install_git.ps1 {} {}"'
               .format(git_link, git_base))
        self._execute(cmd)

    def replace_code(self):
        """Replace the code of cloudbaseinit."""
        LOG.info("Replacing cloudbaseinit's code.")

        LOG.info("Getting program files location.")
        # Get the program files location.
        program_files = self.get_program_files()

        # Remove everything from the cloudbaseinit installation.
        LOG.info("Removing recursively cloudbaseinit.")
        cloudbaseinit = ntpath.join(
            program_files, "Cloudbase Solutions",
            "Cloudbase-Init",
            # TODO(cpopa): take care of this when testing Python 3.
            "Python27",
            "Lib",
            "site-packages",
            "cloudbaseinit")
        self._execute('rmdir "{}" /S /q'.format(cloudbaseinit))

        # Clone the repo
        LOG.info("cloning the cloudbaseinit repo.")
        self._execute("git clone https://github.com/stackforge/"
                      "cloudbase-init C:\\cloudbaseinit")

        # Run the command provided at cli.
        LOG.info("Applying cli patch.")
        opts = util.parse_cli()

        if not opts.git_command:
            raise exceptions.CloudbaseCLIError(
                "git_command not passed at command line. "
                "Could not replace anything.")

        self._execute("cd C:\\cloudbaseinit && {}".format(opts.git_command))

        # Replace the code, by moving the code from cloudbaseinit
        # to the installed location.
        LOG.info("Replacing code.")
        self._execute('powershell "Copy-Item C:\\cloudbaseinit\\cloudbaseinit '
                      '\'{}\' -Recurse"'.format(cloudbaseinit))

    def sysprep(self):
        """Prepare the instance for the actual tests, by running sysprep."""
        LOG.info("Running sysprep.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/sysprep.ps1 -outfile 'C:\\sysprep.ps1'"
               .format(CONF.argus.resources))
        self._execute(cmd)
        self._execute('powershell C:\\sysprep.ps1')

    def wait_cbinit_finalization(self):
        """Wait for the finalization of CloudbaseInit.

        The function waits until all the plugins have been executed.
        """
        LOG.info("Waiting for the finalization of CloudbaseInit execution")

        # Test that this instance's cloudbaseinit run exists.
        key = ('HKLM:SOFTWARE\\Wow6432Node\\Cloudbase '
               'Solutions\\Cloudbase-init\\{0}'
               .format(self._instance_id))
        self._run_cmd_until_condition(
            'powershell Test-Path "{0}"'.format(key),
            lambda out: out.strip() == 'True')

        # Test the number of executed cloudbaseinit plugins.
        wait_cmd = 'powershell (Get-Item %s\\Plugins).ValueCount' % key
        self._run_cmd_until_condition(
            wait_cmd,
            lambda out: int(out) >= int(CONF.argus.expected_plugins_count))

    def wait_reboot(self):
        """Do a reboot and wait until the instance is up."""

        LOG.info('Waiting for server status SHUTOFF because of sysprep')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='SHUTOFF',
            extra_timeout=600)

        self._servers_client.start(self._instance_id)

        LOG.info('Waiting for server status ACTIVE')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='ACTIVE')

<|endoftext|>


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

import abc
import contextlib
import logging
import ntpath
import os
import time

import bs4
import six
from six.moves import urllib

from argus import exceptions
from argus import util

CONF = util.get_config()
LOG = util.get_logger()

__all__ = (
    'InstancePreparer',
    'WindowsInstancePreparer',
)

def _read_url(url):
    request = urllib.request.urlopen(url)
    with contextlib.closing(request) as stream:
        content = stream.read()
        if six.PY3:
            content = content.decode(errors='replace')
        return content

def _get_git_link():
    content = _read_url("http://git-scm.com/download/win")
    soup = bs4.BeautifulSoup(content)
    download_div = soup.find('div', {'class': 'callout downloading'})
    if not download_div:
        raise exceptions.CloudbaseCIError(
            "Could not find callout_downloading div.")

    for a_object in download_div.find_all('a'):
        href = a_object.get('href', '')
        if not href.endswith('.exe'):
            continue
        return href
    raise exceptions.CloudbaseCIError("git download link not found.")

@six.add_metaclass(abc.ABCMeta)
class InstancePreparer(object):
    """Handle instance preparing.

    The method :meth:`~prepare` does all the necessary work for
    preparing a new instance. The executed steps are:

    * wait for boot completion.
    * get an install script for CloudbaseInit
    * installs CloudbaseInit
    * waits for the finalization of the installation.
    """

    def __init__(self, instance_id, servers_client, remote_client):
        self._servers_client = servers_client
        self._instance_id = instance_id
        self._remote_client = remote_client

    def _execute(self, cmd):
        """Execute the given command and fail when the command fails."""
        stdout, stderr, return_code = self._remote_client.run_remote_cmd(cmd)
        if return_code:
            raise exceptions.CloudbaseCIError(
                "Command {command!r} failed with "
                "return code {return_code!r}"
                .format(command=cmd,
                        return_code=return_code))
        return stdout, stderr

    def _run_cmd_until_condition(self, cmd, cond, retry_count=None,
                                 retry_count_interval=5):
        """Run the given `cmd` until a condition *cond* occurs.

        :param cmd:
            A string, representing a command which needs to
            be executed on the underlying remote client.
        :param cond:
            A callable which receives the stdout returned by
            executing the command. It should return a boolean value,
            which tells to this function to stop execution.
        :param retry_count:
            The number of retries which this function has.
            If the value is ``None``, then the function will run *forever*.
        :param retry_count_interval:
            The number of seconds to sleep when retrying a command.
        """
        count = 0
        while True:
            try:
                std_out, std_err = self._execute(cmd)
            except Exception:
                LOG.debug("Command %r failed while waiting for condition",
                          cmd)
                count += 1
                if retry_count and count >= retry_count:
                    raise exceptions.CloudbaseTimeoutError(
                        "Command {!r} failed too many times."
                        .format(cmd))
                time.sleep(retry_count_interval)
            else:
                if std_err:
                    raise exceptions.CloudbaseCLIError(
                        "Executing command {!r} failed with {!r}"
                        .format(cmd, std_err))
                elif cond(std_out):
                    break
                else:
                    time.sleep(retry_count_interval)

    @abc.abstractmethod
    def wait_for_boot_completion(self):
        """Wait for the instance to finish up booting."""

    @abc.abstractmethod
    def get_installation_script(self):
        """Get the installation script for cloudbaseinit."""

    @abc.abstractmethod
    def install_cbinit(self):
        """Install the cloudbaseinit code."""

    @abc.abstractmethod
    def wait_cbinit_finalization(self):
        """Wait for the finalization of cloudbaseinit."""

    @abc.abstractmethod
    def wait_reboot(self):
        """Do a reboot and wait for the instance to be up."""

    @abc.abstractmethod
    def install_git(self):
        """Install git in the instance."""

    @abc.abstractmethod
    def sysprep(self):
        """Do the final steps after installing cloudbaseinit.

        This requires running sysprep on Windows, but on other
        platforms there might be no need for calling it.
        """

    @abc.abstractmethod
    def replace_code(self):
        """Do whatever is necessary to replace the code for cloudbaseinit."""

    def prepare(self):
        """Prepare the underlying instance.

        The following operations will be executed:

        * wait for boot completion
        * get an installation script for CloudbaseInit
        * install CloudbaseInit by running the previously downloaded file.
        * wait until the instance is up and running.
        """
        LOG.info("Preparing instance %s", self._instance_id)
        self.wait_for_boot_completion()
        self.get_installation_script()
        self.install_cbinit()
        self.install_git()

        if CONF.argus.replace_code:
            self.replace_code()

        self.sysprep()
        self.wait_reboot()
        self.wait_cbinit_finalization()
        LOG.info("Finished preparing instance %s", self._instance_id)

    if CONF.argus.debug:
        prepare = util.trap_failure(prepare)

class WindowsInstancePreparer(InstancePreparer):
    """Instance preparer for Windows machines."""

    def get_program_files(self):
        """Get the location of program files from the instance."""
        stdout, _ = self._execute('powershell "(Get-WmiObject  Win32_OperatingSystem).'
                                  'OSArchitecture"')
        architecture = stdout.strip()

        # Next, get the location.
        if architecture == '64-bit':
            location, _ = self._execute('powershell "${ENV:ProgramFiles(x86)}"')
        else:
            location, _ = self._execute('powershell "$ENV:ProgramFiles"')
        return location.strip()

    def wait_for_boot_completion(self):
        LOG.info("Waiting for boot completion")

        wait_cmd = ('powershell "(Get-WmiObject Win32_Account | '
                    'where -Property Name -contains {0}).Name"'
                    .format(CONF.argus.default_ci_username))
        return self._run_cmd_until_condition(
            wait_cmd,
            lambda stdout: stdout.strip() == CONF.argus.default_ci_username)

    def get_installation_script(self):
        """Get an insallation script for CloudbaseInit."""
        LOG.info("Retrieve an installation script for CloudbaseInit")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/installCBinit.ps1 -outfile C:\\installcbinit.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

    def install_cbinit(self):
        """Run the installation script for CloudbaseInit."""
        LOG.info("Run the downloaded installation script")

        cmd = ('powershell "C:\\\\installcbinit.ps1 -serviceType {}"'
               .format(CONF.argus.service_type))
        self._execute(cmd)

    def install_git(self):
        """Install git in the instance."""
        LOG.info("Installing git.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/install_git.ps1 -outfile C:\\\\install_git.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

        git_link = _get_git_link()
        git_base = os.path.basename(git_link)
        cmd = ('powershell "C:\\install_git.ps1 {} {}"'
               .format(git_link, git_base))
        self._execute(cmd)

    def replace_code(self):
        """Replace the code of cloudbaseinit."""
        LOG.info("Replacing cloudbaseinit's code.")

        LOG.info("Getting program files location.")
        # Get the program files location.
        program_files = self.get_program_files()

        # Remove everything from the cloudbaseinit installation.
        LOG.info("Removing recursively cloudbaseinit.")
        cloudbaseinit = ntpath.join(
            program_files, "Cloudbase Solutions",
            "Cloudbase-Init",
            # TODO(cpopa): take care of this when testing Python 3.
            "Python27",
            "Lib",
            "site-packages",
            "cloudbaseinit")
        self._execute('rmdir "{}" /S /q'.format(cloudbaseinit))

        # Clone the repo
        LOG.info("cloning the cloudbaseinit repo.")
        self._execute("git clone https://github.com/stackforge/"
                      "cloudbase-init C:\\cloudbaseinit")

        # Run the command provided at cli.
        LOG.info("Applying cli patch.")
        opts = util.parse_cli()

        if not opts.git_command:
            raise exceptions.CloudbaseCLIError(
                "git_command not passed at command line. "
                "Could not replace anything.")

        self._execute("cd C:\\cloudbaseinit && {}".format(opts.git_command))

        # Replace the code, by moving the code from cloudbaseinit
        # to the installed location.
        LOG.info("Replacing code.")
        self._execute('powershell "Copy-Item C:\\cloudbaseinit\\cloudbaseinit '
                      '\'{}\' -Recurse"'.format(cloudbaseinit))

    def sysprep(self):
        """Prepare the instance for the actual tests, by running sysprep."""
        LOG.info("Running sysprep.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/sysprep.ps1 -outfile 'C:\\sysprep.ps1'"
               .format(CONF.argus.resources))
        self._execute(cmd)
        self._execute('powershell C:\\sysprep.ps1')

    def wait_cbinit_finalization(self):
        """Wait for the finalization of CloudbaseInit.

        The function waits until all the plugins have been executed.
        """
        LOG.info("Waiting for the finalization of CloudbaseInit execution")

        key = ('HKLM:SOFTWARE\\Wow6432Node\\Cloudbase` '
               'Solutions\\Cloudbase-init\\{0}\\Plugins'
               .format(self._instance_id))
        wait_cmd = 'powershell (Get-Item %s).ValueCount' % key

        self._run_cmd_until_condition(
            wait_cmd,
            lambda out: int(out) >= int(CONF.argus.expected_plugins_count))

    def wait_reboot(self):
        """Do a reboot and wait until the instance is up."""

        LOG.info('Waiting for server status SHUTOFF because of sysprep')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='SHUTOFF',
            extra_timeout=600)

        self._servers_client.start(self._instance_id)

        LOG.info('Waiting for server status ACTIVE')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='ACTIVE')

<commit_msg>Wait for the cloudbaseinit's registry key to appear.<commit_after># Copyright 2014 Cloudbase-init
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
"""Instance preparing utilities."""

import abc
import contextlib
import logging
import ntpath
import os
import time

import bs4
import six
from six.moves import urllib

from argus import exceptions
from argus import util

CONF = util.get_config()
LOG = util.get_logger()

__all__ = (
    'InstancePreparer',
    'WindowsInstancePreparer',
)

def _read_url(url):
    request = urllib.request.urlopen(url)
    with contextlib.closing(request) as stream:
        content = stream.read()
        if six.PY3:
            content = content.decode(errors='replace')
        return content

def _get_git_link():
    content = _read_url("http://git-scm.com/download/win")
    soup = bs4.BeautifulSoup(content)
    download_div = soup.find('div', {'class': 'callout downloading'})
    if not download_div:
        raise exceptions.CloudbaseCIError(
            "Could not find callout_downloading div.")

    for a_object in download_div.find_all('a'):
        href = a_object.get('href', '')
        if not href.endswith('.exe'):
            continue
        return href
    raise exceptions.CloudbaseCIError("git download link not found.")

@six.add_metaclass(abc.ABCMeta)
class InstancePreparer(object):
    """Handle instance preparing.

    The method :meth:`~prepare` does all the necessary work for
    preparing a new instance. The executed steps are:

    * wait for boot completion.
    * get an install script for CloudbaseInit
    * installs CloudbaseInit
    * waits for the finalization of the installation.
    """

    def __init__(self, instance_id, servers_client, remote_client):
        self._servers_client = servers_client
        self._instance_id = instance_id
        self._remote_client = remote_client

    def _execute(self, cmd):
        """Execute the given command and fail when the command fails."""
        stdout, stderr, return_code = self._remote_client.run_remote_cmd(cmd)
        if return_code:
            raise exceptions.CloudbaseCIError(
                "Command {command!r} failed with "
                "return code {return_code!r}"
                .format(command=cmd,
                        return_code=return_code))
        return stdout, stderr

    def _run_cmd_until_condition(self, cmd, cond, retry_count=None,
                                 retry_count_interval=5):
        """Run the given `cmd` until a condition *cond* occurs.

        :param cmd:
            A string, representing a command which needs to
            be executed on the underlying remote client.
        :param cond:
            A callable which receives the stdout returned by
            executing the command. It should return a boolean value,
            which tells to this function to stop execution.
        :param retry_count:
            The number of retries which this function has.
            If the value is ``None``, then the function will run *forever*.
        :param retry_count_interval:
            The number of seconds to sleep when retrying a command.
        """
        count = 0
        while True:
            try:
                std_out, std_err = self._execute(cmd)
            except Exception:
                LOG.debug("Command %r failed while waiting for condition",
                          cmd)
                count += 1
                if retry_count and count >= retry_count:
                    raise exceptions.CloudbaseTimeoutError(
                        "Command {!r} failed too many times."
                        .format(cmd))
                time.sleep(retry_count_interval)
            else:
                if std_err:
                    raise exceptions.CloudbaseCLIError(
                        "Executing command {!r} failed with {!r}"
                        .format(cmd, std_err))
                elif cond(std_out):
                    break
                else:
                    time.sleep(retry_count_interval)

    @abc.abstractmethod
    def wait_for_boot_completion(self):
        """Wait for the instance to finish up booting."""

    @abc.abstractmethod
    def get_installation_script(self):
        """Get the installation script for cloudbaseinit."""

    @abc.abstractmethod
    def install_cbinit(self):
        """Install the cloudbaseinit code."""

    @abc.abstractmethod
    def wait_cbinit_finalization(self):
        """Wait for the finalization of cloudbaseinit."""

    @abc.abstractmethod
    def wait_reboot(self):
        """Do a reboot and wait for the instance to be up."""

    @abc.abstractmethod
    def install_git(self):
        """Install git in the instance."""

    @abc.abstractmethod
    def sysprep(self):
        """Do the final steps after installing cloudbaseinit.

        This requires running sysprep on Windows, but on other
        platforms there might be no need for calling it.
        """

    @abc.abstractmethod
    def replace_code(self):
        """Do whatever is necessary to replace the code for cloudbaseinit."""

    def prepare(self):
        """Prepare the underlying instance.

        The following operations will be executed:

        * wait for boot completion
        * get an installation script for CloudbaseInit
        * install CloudbaseInit by running the previously downloaded file.
        * wait until the instance is up and running.
        """
        LOG.info("Preparing instance %s", self._instance_id)
        self.wait_for_boot_completion()
        self.get_installation_script()
        self.install_cbinit()
        self.install_git()

        if CONF.argus.replace_code:
            self.replace_code()

        self.sysprep()
        self.wait_reboot()
        self.wait_cbinit_finalization()
        LOG.info("Finished preparing instance %s", self._instance_id)

    if CONF.argus.debug:
        prepare = util.trap_failure(prepare)

class WindowsInstancePreparer(InstancePreparer):
    """Instance preparer for Windows machines."""

    def get_program_files(self):
        """Get the location of program files from the instance."""
        stdout, _ = self._execute('powershell "(Get-WmiObject  Win32_OperatingSystem).'
                                  'OSArchitecture"')
        architecture = stdout.strip()

        # Next, get the location.
        if architecture == '64-bit':
            location, _ = self._execute('powershell "${ENV:ProgramFiles(x86)}"')
        else:
            location, _ = self._execute('powershell "$ENV:ProgramFiles"')
        return location.strip()

    def wait_for_boot_completion(self):
        LOG.info("Waiting for boot completion")

        wait_cmd = ('powershell "(Get-WmiObject Win32_Account | '
                    'where -Property Name -contains {0}).Name"'
                    .format(CONF.argus.default_ci_username))
        return self._run_cmd_until_condition(
            wait_cmd,
            lambda stdout: stdout.strip() == CONF.argus.default_ci_username)

    def get_installation_script(self):
        """Get an insallation script for CloudbaseInit."""
        LOG.info("Retrieve an installation script for CloudbaseInit")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/installCBinit.ps1 -outfile C:\\installcbinit.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

    def install_cbinit(self):
        """Run the installation script for CloudbaseInit."""
        LOG.info("Run the downloaded installation script")

        cmd = ('powershell "C:\\\\installcbinit.ps1 -serviceType {}"'
               .format(CONF.argus.service_type))
        self._execute(cmd)

    def install_git(self):
        """Install git in the instance."""
        LOG.info("Installing git.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/install_git.ps1 -outfile C:\\\\install_git.ps1"
               .format(CONF.argus.resources))
        self._execute(cmd)

        git_link = _get_git_link()
        git_base = os.path.basename(git_link)
        cmd = ('powershell "C:\\install_git.ps1 {} {}"'
               .format(git_link, git_base))
        self._execute(cmd)

    def replace_code(self):
        """Replace the code of cloudbaseinit."""
        LOG.info("Replacing cloudbaseinit's code.")

        LOG.info("Getting program files location.")
        # Get the program files location.
        program_files = self.get_program_files()

        # Remove everything from the cloudbaseinit installation.
        LOG.info("Removing recursively cloudbaseinit.")
        cloudbaseinit = ntpath.join(
            program_files, "Cloudbase Solutions",
            "Cloudbase-Init",
            # TODO(cpopa): take care of this when testing Python 3.
            "Python27",
            "Lib",
            "site-packages",
            "cloudbaseinit")
        self._execute('rmdir "{}" /S /q'.format(cloudbaseinit))

        # Clone the repo
        LOG.info("cloning the cloudbaseinit repo.")
        self._execute("git clone https://github.com/stackforge/"
                      "cloudbase-init C:\\cloudbaseinit")

        # Run the command provided at cli.
        LOG.info("Applying cli patch.")
        opts = util.parse_cli()

        if not opts.git_command:
            raise exceptions.CloudbaseCLIError(
                "git_command not passed at command line. "
                "Could not replace anything.")

        self._execute("cd C:\\cloudbaseinit && {}".format(opts.git_command))

        # Replace the code, by moving the code from cloudbaseinit
        # to the installed location.
        LOG.info("Replacing code.")
        self._execute('powershell "Copy-Item C:\\cloudbaseinit\\cloudbaseinit '
                      '\'{}\' -Recurse"'.format(cloudbaseinit))

    def sysprep(self):
        """Prepare the instance for the actual tests, by running sysprep."""
        LOG.info("Running sysprep.")

        cmd = ("powershell Invoke-webrequest -uri "
               "{}/sysprep.ps1 -outfile 'C:\\sysprep.ps1'"
               .format(CONF.argus.resources))
        self._execute(cmd)
        self._execute('powershell C:\\sysprep.ps1')

    def wait_cbinit_finalization(self):
        """Wait for the finalization of CloudbaseInit.

        The function waits until all the plugins have been executed.
        """
        LOG.info("Waiting for the finalization of CloudbaseInit execution")

        # Test that this instance's cloudbaseinit run exists.
        key = ('HKLM:SOFTWARE\\Wow6432Node\\Cloudbase` '
               'Solutions\\Cloudbase-init\\{0}'
               .format(self._instance_id))
        self._run_cmd_until_condition(
            'powershell "Test-Path {0}"'.format(key),
            lambda out: out == 'True')

        # Test the number of executed cloudbaseinit plugins.
        wait_cmd = 'powershell (Get-Item %s\\Plugins).ValueCount' % key
        self._run_cmd_until_condition(
            wait_cmd,
            lambda out: int(out) >= int(CONF.argus.expected_plugins_count))

    def wait_reboot(self):
        """Do a reboot and wait until the instance is up."""

        LOG.info('Waiting for server status SHUTOFF because of sysprep')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='SHUTOFF',
            extra_timeout=600)

        self._servers_client.start(self._instance_id)

        LOG.info('Waiting for server status ACTIVE')
        self._servers_client.wait_for_server_status(
            server_id=self._instance_id,
            status='ACTIVE')

<|endoftext|>


<commit_before>#
# ========================================================
# Copyright (c) 2012 Whamcloud, Inc.  All rights reserved.
# ========================================================

import Queue
import json
import traceback
import datetime
import dateutil
import dateutil.tz

from django.db import transaction
from django.http import HttpResponseNotAllowed, HttpResponse, HttpResponseBadRequest
from django.views.decorators.csrf import csrf_exempt
from django.views.generic import View
from functools import wraps
import settings
from tastypie.http import HttpForbidden

from chroma_core.models import ManagedHost
from chroma_core.models.utils import Version
from chroma_core.models.registration_token import RegistrationToken
from chroma_core.services import log_register
from chroma_core.services.http_agent.crypto import Crypto

log = log_register('agent_views')

def log_exception(f):
    @wraps(f)
    def wrapped(*args, **kwargs):
        try:
            return f(*args)
        except Exception:
            log.error(traceback.format_exc())
            raise

    return wrapped

class MessageView(View):
    queues = None
    sessions = None
    hosts = None

    LONG_POLL_TIMEOUT = 30

    @log_exception
    def post(self, request):
        """
        Receive messages FROM the agent.
        Handle a POST containing messages from the agent
        """
        body = json.loads(request.body)
        fqdn = request.META['HTTP_X_SSL_CLIENT_NAME']

        try:
            messages = body['messages']
        except KeyError:
            return HttpResponseBadRequest("Missing attribute 'messages'")

        # Check that the server identifier in each message
        # is valid by comparing against the SSL_CLIENT_NAME
        # which is cryptographically vouched for at the HTTPS frontend
        for message in messages:
            if message['fqdn'] != fqdn:
                return HttpResponseBadRequest("Incorrect client name")

        log.debug("MessageView.post: %s %s messages" % (fqdn, len(messages)))
        for message in messages:
            if message['type'] == 'DATA':
                try:
                    self.sessions.get(fqdn, message['plugin'], message['session_id'])
                except KeyError:
                    log.warning("Terminating session because unknown %s/%s/%s" % (fqdn, message['plugin'], message['session_id']))
                    self.queues.send({
                        'fqdn': fqdn,
                        'type': 'SESSION_TERMINATE',
                        'plugin': message['plugin'],
                        'session_id': None,
                        'session_seq': None,
                        'body': None
                    })
                else:
                    log.debug("Forwarding valid message %s/%s/%s-%s" % (fqdn, message['plugin'], message['session_id'], message['session_seq']))
                    self.queues.receive(message)

            elif message['type'] == 'SESSION_CREATE_REQUEST':
                session = self.sessions.create(fqdn, message['plugin'])
                log.info("Creating session %s/%s/%s" % (fqdn, message['plugin'], session.id))
                self.queues.send({
                    'fqdn': fqdn,
                    'type': 'SESSION_CREATE_RESPONSE',
                    'plugin': session.plugin,
                    'session_id': session.id,
                    'session_seq': None,
                    'body': None
                })

        return HttpResponse()

    def _filter_valid_messages(self, fqdn, messages):
        plugin_to_session_id = {}

        def is_valid(message):
            try:
                session_id = plugin_to_session_id[message['plugin']]
            except KeyError:
                try:
                    plugin_to_session_id[message['plugin']] = session_id = self.sessions.get(fqdn, message['plugin']).id
                except KeyError:
                    plugin_to_session_id[message['plugin']] = session_id = None

            if message['session_id'] != session_id:
                log.debug("Dropping message because it has stale session id (current is %s): %s" % (session_id, message))
                return False

            return True

        return [m for m in messages if is_valid(m)]

    @log_exception
    def get(self, request):
        """
        Send messages TO the agent.
        Handle a long-polling GET for messages to the agent
        """

        fqdn = request.META['HTTP_X_SSL_CLIENT_NAME']
        server_boot_time = dateutil.parser.parse(request.GET['server_boot_time'])
        client_start_time = dateutil.parser.parse(request.GET['client_start_time'])

        messages = []

        try:
            reset_required = self.hosts.update(fqdn, server_boot_time, client_start_time)
        except ManagedHost.DoesNotExist:
            # This should not happen because the HTTPS frontend should have the
            # agent certificate revoked before removing the ManagedHost from the database
            log.error("GET from unknown server %s" % fqdn)
            return HttpResponseBadRequest("Unknown server '%s'" % fqdn)

        if reset_required:
            # This is the case where the http_agent service restarts, so
            # we have to let the agent know that all open sessions
            # are now over.
            messages.append({
                'fqdn': fqdn,
                'type': 'SESSION_TERMINATE_ALL',
                'plugin': None,
                'session_id': None,
                'session_seq': None,
                'body': None
            })

        log.info("MessageView.get: composing messages for %s" % fqdn)
        queue = self.queues.get(fqdn).tx

        try:
            first_message = queue.get(block = True, timeout = self.LONG_POLL_TIMEOUT)
        except Queue.Empty:
            pass
        else:
            # TODO: limit number of messages per response
            messages.append(first_message)
            while True:
                try:
                    messages.append(queue.get(block = False))
                except Queue.Empty:
                    break

        messages = self._filter_valid_messages(fqdn, messages)

        log.info("MessageView.get: responding to %s with %s messages" % (fqdn, len(messages)))
        return HttpResponse(json.dumps({'messages': messages}), mimetype = "application/json")

@csrf_exempt
@log_exception
def register(request, key):
    if request.method != "POST":
        return HttpResponseNotAllowed(['POST'])

    # Validate the secret key
    try:
        with transaction.commit_on_success():
            token = RegistrationToken.objects.get(secret = key)
            if not token.credits:
                log.warning("Attempt to register with exhausted token %s" % key)
                return HttpForbidden()
            else:
                # Decrement .credits
                RegistrationToken.objects.filter(secret = key).update(credits = token.credits - 1)
    except RegistrationToken.DoesNotExist:
        log.warning("Attempt to register with non-existent token %s" % key)
        return HttpForbidden()
    else:
        now = datetime.datetime.utcnow()
        now = now.replace(tzinfo = dateutil.tz.tzutc())
        log.info("now.tzinfo= %s" % now.tzinfo)
        log.info("expiry.tzinfo = %s" % token.expiry)

        if token.expiry < now:
            log.warning("Attempt to register with expired token %s (now %s, expired at %s)" % (key, now, token.expiry))
            return HttpForbidden()
        elif token.cancelled:
            log.warning("Attempt to register with cancelled token %s" % key)
            return HttpForbidden()

    host_attributes = json.loads(request.body)

    # Fail at the first if the version of the agent on the server is incorrect
    manager, agent = Version(settings.VERSION), Version(host_attributes['version'])
    if manager and agent and not (manager.major == agent.major and manager.minor >= agent.minor):
        return HttpResponse(status = 400, content = "Version incompatibility between manager {0} and agent {1}".format(manager, agent))

    # Fulfil the registering server's request for a certificate authenticating
    # it as the owner of this FQDN.
    csr = host_attributes['csr']

    # Check that the commonName in the CSR is the same as that in host_attributes
    # (prevent registering as one host and getting a certificate to impersonate another)
    csr_fqdn = Crypto().get_common_name(csr)
    if csr_fqdn != host_attributes['fqdn']:
        # Terse response to attacker
        log.error("FQDN mismatch '%s' vs. '%s' from %s" % (csr_fqdn, host_attributes['fqdn'], request.META['HTTP_X_FORWARDED_FOR']))
        return HttpResponse(status = 400, content = "")

    certificate_str = Crypto().sign(csr)

    # FIXME: handle the case where someone registers,
    # and then dies before saving their certificate:
    # when they come through here again, currently
    # we'll reject them because the FQDN is taken
    # ... maybe hand back the certificate here, but
    # then don't create the host until they first
    # connect using the certificate?
    # in that case to avoid handing out another cert
    # to someone else spamming our URL, we should have
    # some logic during the second addition to revoke
    # the first (should never be used) host cert.

    host, command = ManagedHost.create(
        address = host_attributes['address'],
        fqdn = host_attributes['fqdn'],
        nodename = host_attributes['nodename'],
        capabilities = host_attributes['capabilities']
    )

    # TODO: document this return format
    return HttpResponse(status = 201, content = json.dumps({
        'command_id': command.id,
        'host_id': host.id,
        'certificate': certificate_str
    }), mimetype = "application/json")
<commit_msg>Tone down logging slightly<commit_after>#
# ========================================================
# Copyright (c) 2012 Whamcloud, Inc.  All rights reserved.
# ========================================================

import Queue
import json
import traceback
import datetime
import dateutil
import dateutil.tz

from django.db import transaction
from django.http import HttpResponseNotAllowed, HttpResponse, HttpResponseBadRequest
from django.views.decorators.csrf import csrf_exempt
from django.views.generic import View
from functools import wraps
import settings
from tastypie.http import HttpForbidden

from chroma_core.models import ManagedHost
from chroma_core.models.utils import Version
from chroma_core.models.registration_token import RegistrationToken
from chroma_core.services import log_register
from chroma_core.services.http_agent.crypto import Crypto

log = log_register('agent_views')

def log_exception(f):
    @wraps(f)
    def wrapped(*args, **kwargs):
        try:
            return f(*args)
        except Exception:
            log.error(traceback.format_exc())
            raise

    return wrapped

class MessageView(View):
    queues = None
    sessions = None
    hosts = None

    LONG_POLL_TIMEOUT = 30

    @log_exception
    def post(self, request):
        """
        Receive messages FROM the agent.
        Handle a POST containing messages from the agent
        """
        body = json.loads(request.body)
        fqdn = request.META['HTTP_X_SSL_CLIENT_NAME']

        try:
            messages = body['messages']
        except KeyError:
            return HttpResponseBadRequest("Missing attribute 'messages'")

        # Check that the server identifier in each message
        # is valid by comparing against the SSL_CLIENT_NAME
        # which is cryptographically vouched for at the HTTPS frontend
        for message in messages:
            if message['fqdn'] != fqdn:
                return HttpResponseBadRequest("Incorrect client name")

        log.debug("MessageView.post: %s %s messages" % (fqdn, len(messages)))
        for message in messages:
            if message['type'] == 'DATA':
                try:
                    self.sessions.get(fqdn, message['plugin'], message['session_id'])
                except KeyError:
                    log.warning("Terminating session because unknown %s/%s/%s" % (fqdn, message['plugin'], message['session_id']))
                    self.queues.send({
                        'fqdn': fqdn,
                        'type': 'SESSION_TERMINATE',
                        'plugin': message['plugin'],
                        'session_id': None,
                        'session_seq': None,
                        'body': None
                    })
                else:
                    log.debug("Forwarding valid message %s/%s/%s-%s" % (fqdn, message['plugin'], message['session_id'], message['session_seq']))
                    self.queues.receive(message)

            elif message['type'] == 'SESSION_CREATE_REQUEST':
                session = self.sessions.create(fqdn, message['plugin'])
                log.info("Creating session %s/%s/%s" % (fqdn, message['plugin'], session.id))
                self.queues.send({
                    'fqdn': fqdn,
                    'type': 'SESSION_CREATE_RESPONSE',
                    'plugin': session.plugin,
                    'session_id': session.id,
                    'session_seq': None,
                    'body': None
                })

        return HttpResponse()

    def _filter_valid_messages(self, fqdn, messages):
        plugin_to_session_id = {}

        def is_valid(message):
            try:
                session_id = plugin_to_session_id[message['plugin']]
            except KeyError:
                try:
                    plugin_to_session_id[message['plugin']] = session_id = self.sessions.get(fqdn, message['plugin']).id
                except KeyError:
                    plugin_to_session_id[message['plugin']] = session_id = None

            if message['session_id'] != session_id:
                log.debug("Dropping message because it has stale session id (current is %s): %s" % (session_id, message))
                return False

            return True

        return [m for m in messages if is_valid(m)]

    @log_exception
    def get(self, request):
        """
        Send messages TO the agent.
        Handle a long-polling GET for messages to the agent
        """

        fqdn = request.META['HTTP_X_SSL_CLIENT_NAME']
        server_boot_time = dateutil.parser.parse(request.GET['server_boot_time'])
        client_start_time = dateutil.parser.parse(request.GET['client_start_time'])

        messages = []

        try:
            reset_required = self.hosts.update(fqdn, server_boot_time, client_start_time)
        except ManagedHost.DoesNotExist:
            # This should not happen because the HTTPS frontend should have the
            # agent certificate revoked before removing the ManagedHost from the database
            log.error("GET from unknown server %s" % fqdn)
            return HttpResponseBadRequest("Unknown server '%s'" % fqdn)

        if reset_required:
            # This is the case where the http_agent service restarts, so
            # we have to let the agent know that all open sessions
            # are now over.
            messages.append({
                'fqdn': fqdn,
                'type': 'SESSION_TERMINATE_ALL',
                'plugin': None,
                'session_id': None,
                'session_seq': None,
                'body': None
            })

        log.debug("MessageView.get: composing messages for %s" % fqdn)
        queue = self.queues.get(fqdn).tx

        try:
            first_message = queue.get(block = True, timeout = self.LONG_POLL_TIMEOUT)
        except Queue.Empty:
            pass
        else:
            # TODO: limit number of messages per response
            messages.append(first_message)
            while True:
                try:
                    messages.append(queue.get(block = False))
                except Queue.Empty:
                    break

        messages = self._filter_valid_messages(fqdn, messages)

        log.debug("MessageView.get: responding to %s with %s messages" % (fqdn, len(messages)))
        return HttpResponse(json.dumps({'messages': messages}), mimetype = "application/json")

@csrf_exempt
@log_exception
def register(request, key):
    if request.method != "POST":
        return HttpResponseNotAllowed(['POST'])

    # Validate the secret key
    try:
        with transaction.commit_on_success():
            token = RegistrationToken.objects.get(secret = key)
            if not token.credits:
                log.warning("Attempt to register with exhausted token %s" % key)
                return HttpForbidden()
            else:
                # Decrement .credits
                RegistrationToken.objects.filter(secret = key).update(credits = token.credits - 1)
    except RegistrationToken.DoesNotExist:
        log.warning("Attempt to register with non-existent token %s" % key)
        return HttpForbidden()
    else:
        now = datetime.datetime.utcnow()
        now = now.replace(tzinfo = dateutil.tz.tzutc())

        if token.expiry < now:
            log.warning("Attempt to register with expired token %s (now %s, expired at %s)" % (key, now, token.expiry))
            return HttpForbidden()
        elif token.cancelled:
            log.warning("Attempt to register with cancelled token %s" % key)
            return HttpForbidden()

    host_attributes = json.loads(request.body)

    # Fail at the first if the version of the agent on the server is incorrect
    manager, agent = Version(settings.VERSION), Version(host_attributes['version'])
    if manager and agent and not (manager.major == agent.major and manager.minor >= agent.minor):
        return HttpResponse(status = 400, content = "Version incompatibility between manager {0} and agent {1}".format(manager, agent))

    # Fulfil the registering server's request for a certificate authenticating
    # it as the owner of this FQDN.
    csr = host_attributes['csr']

    # Check that the commonName in the CSR is the same as that in host_attributes
    # (prevent registering as one host and getting a certificate to impersonate another)
    csr_fqdn = Crypto().get_common_name(csr)
    if csr_fqdn != host_attributes['fqdn']:
        # Terse response to attacker
        log.error("FQDN mismatch '%s' vs. '%s' from %s" % (csr_fqdn, host_attributes['fqdn'], request.META['HTTP_X_FORWARDED_FOR']))
        return HttpResponse(status = 400, content = "")

    certificate_str = Crypto().sign(csr)

    # FIXME: handle the case where someone registers,
    # and then dies before saving their certificate:
    # when they come through here again, currently
    # we'll reject them because the FQDN is taken
    # ... maybe hand back the certificate here, but
    # then don't create the host until they first
    # connect using the certificate?
    # in that case to avoid handing out another cert
    # to someone else spamming our URL, we should have
    # some logic during the second addition to revoke
    # the first (should never be used) host cert.

    host, command = ManagedHost.create(
        address = host_attributes['address'],
        fqdn = host_attributes['fqdn'],
        nodename = host_attributes['nodename'],
        capabilities = host_attributes['capabilities']
    )

    # TODO: document this return format
    return HttpResponse(status = 201, content = json.dumps({
        'command_id': command.id,
        'host_id': host.id,
        'certificate': certificate_str
    }), mimetype = "application/json")
<|endoftext|>


<commit_before>from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.contrib import messages
from django.http import HttpResponseForbidden
from django.shortcuts import redirect, render_to_response, get_object_or_404
from django.views.generic.create_update import update_object
from django.views.generic.list_detail import object_list, object_detail
from django.utils.translation import ugettext as _
from django.views.decorators.http import require_GET

from haystack.query import SearchQuerySet
import notification.models as notification
from voting.views import vote_on_object

from .models import Feedback
from .forms import FeedbackNewForm, FeedbackEditForm
from .signals import feedback_updated as sig_feedback_updated
from .sql import SumWithDefault

@login_required
def feedback_new(request, template_name='backcap/feedback_new.html'):
    """
    Create a new feedback
    """
    referer = request.GET.get("referer", None)

    if request.method == 'POST':
        feedback_form = FeedbackNewForm(request.POST)

        if feedback_form.is_valid():
            feedback = feedback_form.save(commit=False)
            feedback.user = request.user
            feedback.save()

            messages.success(request, _("Thanks for you feedback !"))

            staff = User.objects.filter(is_staff=True)
            notification.send(staff, "feedback_new", {'feedback': feedback})
            notification.observe(feedback, request.user, "feedback_updated", "feedback_updated")

            return redirect(feedback)
    else:
        feedback_form = FeedbackNewForm(initial={'referer': referer})

    return render_to_response(template_name=template_name,
                              dictionary={'feedback_form': feedback_form},
                              )

# XXX: Security
@login_required
def feedback_update(request, feedback_id):
    """
    Edit a single feedback
    """
    return update_object(request,
                         form_class=FeedbackEditForm,
                         object_id=feedback_id,
                         template_name='backcap/feedback_update.html',
                         )

def feedback_list(request, qtype='all'):
    """
    Display all the feedbacks
    """
    queryset = Feedback.objects.exclude(status__in=('C', 'D', 'I')).annotate(score=SumWithDefault('votes__vote', default=0))

    order = request.GET.get('order', 'score')
    if order == 'newest':
        quersyet = queryset.order_by('-modified_on', 'kind', '-score')
    else:
        queryset = queryset.order_by('-score', '-modified_on', 'kind')

    if qtype in [choice[0] for choice in Feedback.KIND_CHOICES]:
        queryset = queryset.filter(kind=qtype)

    return object_list(request,
                       queryset=queryset,
                       template_name='backcap/feedback_list.html',
                       template_object_name='feedback',
                       paginate_by=15,
                       extra_context={'qtype': qtype,
                                      'order': order},
                       )

def feedback_detail(request, feedback_id):
    """
    Shows a single feedback
    """
    return object_detail(request,
                         queryset=Feedback.objects.all(),
                         object_id=feedback_id,
                         template_object_name='feedback',
                         template_name='backcap/feedback_detail.html',
                         )

# XXX Security
@login_required
def feedback_close(request, feedback_id):
    """
    Closes a feedback. This means it has been resolved.
    """
    feedback = get_object_or_404(Feedback, id=feedback_id)

    # Close it, then save
    feedback.status = 'C'
    feedback.save()

    return redirect(feedback)

@login_required
def feedback_vote(request, feedback_id, direction):
    feedback = get_object_or_404(Feedback, id=feedback_id)

    # Auto (un)subscribe user if he/she's interested in this issue or not
    is_observing = notification.is_observing(feedback, request.user, "feedback_updated")
    if direction == 'up' and not is_observing:
        notification.observe(feedback, request.user, "feedback_updated", "feedback_updated")
    elif direction in ('down', 'clear') and is_observing:
        notification.stop_observing(feedback, request.user, "feedback_updated")

    return vote_on_object(request,
                          model=Feedback,
                          direction=direction,
                          object_id=feedback_id,
                          template_object_name='vote',
                          template_name='kb/link_confirm_vote.html',
                          allow_xmlhttprequest=True)

@login_required
def feedback_ping_observers(request, feedback_id):
    """
    Ping users that are observing this feedback
    """
    if not request.user.is_superuser:
        return HttpResponseForbidden()

    feedback = get_object_or_404(Feedback, id=feedback_id)

    notification.send([feedback.user],
                      "feedback_updated",
                      {'feedback': feedback},
                      )

    messages.success(request, _("The users have been successfully notified"))

    return redirect(feedback)

def feedback_tab(request):
    return feedback_new(request,
                        template_name='backcap/feedback_tab.html')

@require_GET
def feedback_search(request, limit=10):
    query = request.GET['q']

    results = SearchQuerySet().models(Feedback).filter_or(content=query)
    if limit:
        results = results[:int(limit)]

    return render_to_response(template_name='backcap/feedback_search.html',
                              dictionary={'results': results,
                                          'query': query},
                              )
<commit_msg>change ordering<commit_after>from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.contrib import messages
from django.http import HttpResponseForbidden
from django.shortcuts import redirect, render_to_response, get_object_or_404
from django.views.generic.create_update import update_object
from django.views.generic.list_detail import object_list, object_detail
from django.utils.translation import ugettext as _
from django.views.decorators.http import require_GET

from haystack.query import SearchQuerySet
import notification.models as notification
from voting.views import vote_on_object

from .models import Feedback
from .forms import FeedbackNewForm, FeedbackEditForm
from .signals import feedback_updated as sig_feedback_updated
from .sql import SumWithDefault

@login_required
def feedback_new(request, template_name='backcap/feedback_new.html'):
    """
    Create a new feedback
    """
    referer = request.GET.get("referer", None)

    if request.method == 'POST':
        feedback_form = FeedbackNewForm(request.POST)

        if feedback_form.is_valid():
            feedback = feedback_form.save(commit=False)
            feedback.user = request.user
            feedback.save()

            messages.success(request, _("Thanks for you feedback !"))

            staff = User.objects.filter(is_staff=True)
            notification.send(staff, "feedback_new", {'feedback': feedback})
            notification.observe(feedback, request.user, "feedback_updated", "feedback_updated")

            return redirect(feedback)
    else:
        feedback_form = FeedbackNewForm(initial={'referer': referer})

    return render_to_response(template_name=template_name,
                              dictionary={'feedback_form': feedback_form},
                              )

# XXX: Security
@login_required
def feedback_update(request, feedback_id):
    """
    Edit a single feedback
    """
    return update_object(request,
                         form_class=FeedbackEditForm,
                         object_id=feedback_id,
                         template_name='backcap/feedback_update.html',
                         )

def feedback_list(request, qtype='all'):
    """
    Display all the feedbacks
    """
    queryset = Feedback.objects.exclude(status__in=('C', 'D', 'I')).annotate(score=SumWithDefault('votes__vote', default=0))

    order = request.GET.get('order', 'score')
    if order == 'newest':
        quersyet = queryset.order_by('modified_on', 'kind', '-score')
    else:
        queryset = queryset.order_by('-score', 'modified_on', 'kind')

    if qtype in [choice[0] for choice in Feedback.KIND_CHOICES]:
        queryset = queryset.filter(kind=qtype)

    return object_list(request,
                       queryset=queryset,
                       template_name='backcap/feedback_list.html',
                       template_object_name='feedback',
                       paginate_by=15,
                       extra_context={'qtype': qtype,
                                      'order': order},
                       )

def feedback_detail(request, feedback_id):
    """
    Shows a single feedback
    """
    return object_detail(request,
                         queryset=Feedback.objects.all(),
                         object_id=feedback_id,
                         template_object_name='feedback',
                         template_name='backcap/feedback_detail.html',
                         )

# XXX Security
@login_required
def feedback_close(request, feedback_id):
    """
    Closes a feedback. This means it has been resolved.
    """
    feedback = get_object_or_404(Feedback, id=feedback_id)

    # Close it, then save
    feedback.status = 'C'
    feedback.save()

    return redirect(feedback)

@login_required
def feedback_vote(request, feedback_id, direction):
    feedback = get_object_or_404(Feedback, id=feedback_id)

    # Auto (un)subscribe user if he/she's interested in this issue or not
    is_observing = notification.is_observing(feedback, request.user, "feedback_updated")
    if direction == 'up' and not is_observing:
        notification.observe(feedback, request.user, "feedback_updated", "feedback_updated")
    elif direction in ('down', 'clear') and is_observing:
        notification.stop_observing(feedback, request.user, "feedback_updated")

    return vote_on_object(request,
                          model=Feedback,
                          direction=direction,
                          object_id=feedback_id,
                          template_object_name='vote',
                          template_name='kb/link_confirm_vote.html',
                          allow_xmlhttprequest=True)

@login_required
def feedback_ping_observers(request, feedback_id):
    """
    Ping users that are observing this feedback
    """
    if not request.user.is_superuser:
        return HttpResponseForbidden()

    feedback = get_object_or_404(Feedback, id=feedback_id)

    notification.send([feedback.user],
                      "feedback_updated",
                      {'feedback': feedback},
                      )

    messages.success(request, _("The users have been successfully notified"))

    return redirect(feedback)

def feedback_tab(request):
    return feedback_new(request,
                        template_name='backcap/feedback_tab.html')

@require_GET
def feedback_search(request, limit=10):
    query = request.GET['q']

    results = SearchQuerySet().models(Feedback).filter_or(content=query)
    if limit:
        results = results[:int(limit)]

    return render_to_response(template_name='backcap/feedback_search.html',
                              dictionary={'results': results,
                                          'query': query},
                              )
<|endoftext|>