<commit_before>#
# Copyright (c) 2011 rPath, Inc.
#

from lxml import etree
from lxml.builder import ElementMaker

class Survey(object):
    """
    Class for adding any additional information to a Windows survey.
    """

    def __init__(self, survey_data, rtis):
        # survey_data is a etree element
        self._data = survey_data
        self._rtis = rtis

        self.e = ElementMaker()

    def tostring(self, prettyPrint=False):
        root = self.e.surveys()
        root.append(self._data)
        return etree.tostring(root, pretty_print=prettyPrint)

    def addPackageInformation(self):
        """
        Add the conary package information to the survey.
        """

        # 1. get the conary package data from the target system
        # 2. get troveInfo for all packages from teh repository
        # 3. map conary packages to windows packages and vice versa

<commit_msg>add conary package info for windows surveys (RCE-55)<commit_after>#
# Copyright (c) 2011 rPath, Inc.
#

import copy
from xml.etree import cElementTree as et

from lxml import etree
from lxml.builder import ElementMaker

from rpath_tools.client.sysdisco.packages import IDFactory
from rpath_tools.client.sysdisco.packages import ConaryScanner as _ConaryScanner
from rpath_tools.client.sysdisco.packages import WindowsScanner as _WindowsScanner


def children(node):
    return dict((x.tag, x) for x in node.iterchildren())


class DbShim(object):
    def __init__(self, repos, pkglist):
        self.repos = repos
        self.pkglist = pkglist

        self.getTrove = self.repos.getTrove

    def iterAllTroves(self):
        return iter(self.pkglist)


class ConaryScanner(_ConaryScanner):
    def __init__(self, pkglist, client):
        _ConaryScanner.__init__(self)
        self.pkglist = pkglist
        self.client = client

    def _getDb(self):
        return DbShim(self.client.repos, self.pkglist)


class WindowsPackage(object):
    __slots__ = ('node', 'pkgInfo', )

    def __init__(self, node):
        self.node = node
        self.pkgInfo = children(self.node).get('windows_package_info')

    @property
    def id(self):
        return self.node.attrib['id']

    @property
    def name(self):
        return children(self.pkgInfo).get('product_name').text

    @property
    def version(self):
        return children(self.pkgInfo).get('version').text

    @property
    def type(self):
        return children(self.pkgInfo).get('type').text

    @property
    def productCode(self):
        return children(self.pkgInfo).get('product_code').text

    def toxml(self):
        return self.node

    def __hash__(self):
        return hash(self.productCode)

    def __cmp__(self, other):
        assert isinstance(other, WindowsPackage)

        code = cmp(self.productCode, other.productCode)
        if code != 0:
            return code

        return cmp(self.name, other.name)


class WindowsScanner(_WindowsScanner):
    def __init__(self, packages):
        _WindowsScanner.__init__(self)
        self.packages = packages

    def scan(self):
        self._results = {}
        for node in self.packages.iterchildren():
            pkg = WindowsPackage(node)
            self._results[pkg.productCode] = pkg
        return self._results


class Survey(object):
    """
    Class for adding any additional information to a Windows survey.
    """

    def __init__(self, survey_data, rtis):
        # survey_data is a etree element
        self._data = survey_data
        self._rtis = rtis

        self.e = ElementMaker()

    def tostring(self, prettyPrint=False):
        root = self.e.surveys()
        root.append(self._data)
        return etree.tostring(root, pretty_print=prettyPrint)

    def _getConaryClient(self):
        # Import here to avoid import loop
        from rpath_repeater.utils.windows.updates import UpdateJob
        return UpdateJob(self._rtis.flavor, self._rtis.manifest,
            None, callback=self._rtis.callback)._client

    def addPackageInformation(self):
        """
        Add the conary package information to the survey.
        """

        # 1. get the conary package data from the target system
        # 2. get troveInfo for all packages from teh repository
        # 3. map conary packages to windows packages and vice versa

        manifest = self._rtis.manifest
        conaryInfo = ConaryScanner(manifest, self._getConaryClient()).scan()
        productCodes = dict((x.msi.productCode, x)
            for x in conaryInfo.itervalues() if x.msi)

        windowsPkgs = children(self._data).get('windows_packages')
        windowsInfo = WindowsScanner(windowsPkgs).scan()

        idGen = IDFactory()

        conary_packages = self.e.conary_packages()
        for pkg in conaryInfo.itervalues():
            nodeId = idGen.getId(pkg)
            node = etree.fromstring(et.tostring(pkg.toxml(nodeId)))
            if pkg.msi:
                winPkg = windowsInfo[pkg.msi.productCode]
                child = node.find('.//conary_package_info')
                child.append(copy.copy(winPkg.node))
            conary_packages.append(node)

        windows_packages = self.e.windows_packages()
        for pkg in windowsInfo.itervalues():
            node = pkg.node
            cnyPkg = productCodes.get(pkg.productCode)
            if cnyPkg:
                node.append(self.e.conary_package(id=idGen.getId(cnyPkg)))
            node.append(self.e.encapsulated(str(bool(cnyPkg)).lower()))
            windows_packages.append(node)

        self._data.remove(children(self._data).get('windows_packages'))
        self._data.append(conary_packages)
        self._data.append(windows_packages)

{'fraction_non_alphanumeric': 0.09560769095078497, 'fraction_numerical': 0.002998765214323514, 'mean_word_length': 2.586337760910816, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 7}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1544586', 'n_tokens_mistral': 1727, 'n_tokens_neox': 1657, 'n_words': 406}