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

import sys

from . import crypto
from . import db

class Election:

    def __init__(self, eid, db_fname):
        self.eid = eid
        self.db = db.DB(db_fname)

        # Construct cursors for all operations.
        self.c_salt_issue = self.db.add_statement(
            '''UPDATE ISSUES SET salt = ? WHERE _ROWID_ = ?''')
        self.c_salt_record = self.db.add_statement(
            '''UPDATE RECORD SET salt = ? WHERE _ROWID_ = ?''')
        self.c_open = self.db.add_statement(
            'UPDATE METADATA SET salt = ?, opened_key = ?')
        self.c_close = self.db.add_statement(
            'UPDATE METADATA SET closed = 1')
        self.c_add_issue = self.db.add_statement(
            '''INSERT INTO ISSUES VALUES (?, ?, ?, ?, ?, ?)
               ON CONFLICT DO UPDATE SET
                 title=excluded.title,
                 description=excluded.description,
                 type=excluded.type,
                 kv=excluded.kv
            ''')
        self.c_add_record = self.db.add_statement(
            '''INSERT INTO RECORD VALUES (?, ?, ?, ?)
               ON CONFLICT DO UPDATE SET
                 name=excluded.name,
                 email=excluded.email
            ''')
        self.c_delete_issue = self.db.add_statement(
            'DELETE FROM ISSUES WHERE iid = ?')
        self.c_delete_record = self.db.add_statement(
            'DELETE FROM RECORD WHERE rid = ?')

        # Cursors for running queries.
        self.q_metadata = self.db.add_query('metadata',
            'SELECT * FROM METADATA')
        self.q_issues = self.db.add_query('issues',
            'SELECT * FROM ISSUES ORDER BY iid')
        self.q_record = self.db.add_query('record',
            'SELECT * FROM RECORD ORDER BY rid')
        self.q_get_issue = self.db.add_query('issues',
            'SELECT * FROM ISSUES WHERE iid = ?')
        self.q_get_record = self.db.add_query('record',
            'SELECT * FROM RECORD WHERE rid = ?')

    def open(self):

        # Double-check the Election is in the editing state.
        assert self.is_editable()

        # Add salts first. If this is gonna fail, then make sure it
        # happens before we move to the "opened" state.
        self.add_salts()

        edata = self.gather_election_data()
        print('EDATA:', edata)
        salt = crypto.gen_salt()
        opened_key = crypto.gen_opened_key(edata, salt)

        print('SALT:', salt)
        print('KEY:', opened_key)
        self.c_open.perform((salt, opened_key))

    def gather_election_data(self):
        "Gather a definition of this election for keying and anti-tamper."

        # NOTE: separators and other zero-entropy constant chars are
        # not included when assembling the data for hashing. This data
        # is not intended for human consumption, anyways.

        # NOTE: all assembly of rows must use a repeatable ordering.

        md = self.q_metadata.first_row()
        mdata = md.eid + md.title

        self.q_issues.perform()
        # Use an f-string to render "None" if a column is NULL.
        idata = ''.join(f'{r.iid}{r.title}{r.description}{r.type}{r.kv}'
                        for r in self.q_issues.fetchall())

        self.q_record.perform()
        rdata = ''.join(r.rid + r.email
                        for r in self.q_record.fetchall())

        return (mdata + idata + rdata).encode()

    def close(self):
        "Close an election."

        # The Election should be open.
        assert self.is_open()

        # Simple tweak of the metadata to close the Election.
        self.c_close.perform()

    def add_salts(self):
        "Set the SALT column in the ISSUES and RECORD tables."

        # The Election should be editable.
        assert self.is_editable()

        cur = self.db.conn.cursor()

        def for_table(table, mod_cursor):
            "Use MOD_CURSOR to salt each row of TABLE."

            # Fetch all ROWID values now, to avoid two cursors
            # attempting to work on TABLE at the same time.
            cur.execute(f'SELECT _ROWID_ FROM {table}')
            ids = cur.fetchall()

            # Now, add a salt to every row.
            for r in ids:
                salt = crypto.gen_salt()
                print('ROW will use:', table, r[0], salt)
                mod_cursor.perform((salt, r[0]))

        for_table('issues', self.c_salt_issue)
        for_table('record', self.c_salt_record)

    def get_issue(self, iid):
        "Return TITLE, DESCRIPTION, TYPE, and KV for issue IID."

        # NEVER return issue.salt
        issue = self.q_get_issue.first_row((iid,))
        return issue.title, issue.description, issue.type, issue.kv

    def add_issue(self, iid, title, description, type, kv):
        "Add or update an issue designated by IID."
        assert self.is_editable()

        # If we ADD, then SALT will be NULL. If we UPDATE, then it will not
        # be touched (it should be NULL).
        self.c_add_issue.perform((iid, title, description, type, kv, None))

    def delete_issue(self, iid):
        "Delete the Issue designated by IID."
        assert self.is_editable()

        self.c_delete_issue.perform((iid,))

    def get_participant(self, rid):
        "Return NAME, EMAIL for Participant on record RID."

        # NEVER return record.salt
        record = self.q_get_record.first_row((rid,))
        return record.name, record.email

    def add_participant(self, rid, name, email):
        "Add or update a Participant (voter of record) designated by RID."
        assert self.is_editable()

        # If we ADD, then SALT will be NULL. If we UPDATE, then it will not
        # be touched (it should be NULL).
        self.c_add_record.perform((rid, name, email, None))

    def delete_participant(self, rid):
        "Delete the Participant designated by RID."
        assert self.is_editable()

        self.c_delete_record.perform((rid,))

    def is_tampered(self):

        # The Election should be open.
        assert self.is_open()

        md = self.q_metadata.first_row()

        # Compute an opened_key based on the current data.
        edata = self.gather_election_data()
        opened_key = crypto.gen_opened_key(edata, md.salt)

        print('EDATA:', edata)
        print('SALT:', md.salt)
        print('KEY:', opened_key)

        # The computed key should be unchanged.
        return opened_key != md.opened_key

    def is_editable(self):
        "Can this Election be edited?"
        md = self.q_metadata.first_row()
        return md.salt is None and md.opened_key is None

    def is_open(self):
        "Is this Election open for voting?"
        md = self.q_metadata.first_row()
        return (md.salt is not None
                and md.opened_key is not None
                and md.closed in (None, 0))

    def is_closed(self):
        "Has this Election been closed?"
        md = self.q_metadata.first_row()
        return (md.salt is None
                and md.opened_key is None
                and md.closed == 1)

def new_eid():
    "Create a new ElectionID."

    # Use 4 bytes of a salt, for 32 bits.
    b = crypto.gen_salt()

    # Format into 8 hex characters.
    return f'{b[0]:02x}{b[1]:02x}{b[2]:02x}{b[3]:02x}'
<commit_msg>implement recording of votes<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# ----
#
# ### TBD: DOCCO
#
#

import sys

from . import crypto
from . import db

class Election:

    def __init__(self, eid, db_fname):
        self.eid = eid
        self.db = db.DB(db_fname)

        # Construct cursors for all operations.
        self.c_salt_issue = self.db.add_statement(
            '''UPDATE ISSUES SET salt = ? WHERE _ROWID_ = ?''')
        self.c_salt_record = self.db.add_statement(
            '''UPDATE RECORD SET salt = ? WHERE _ROWID_ = ?''')
        self.c_open = self.db.add_statement(
            'UPDATE METADATA SET salt = ?, opened_key = ?')
        self.c_close = self.db.add_statement(
            'UPDATE METADATA SET closed = 1')
        self.c_add_issue = self.db.add_statement(
            '''INSERT INTO ISSUES VALUES (?, ?, ?, ?, ?, ?)
               ON CONFLICT DO UPDATE SET
                 title=excluded.title,
                 description=excluded.description,
                 type=excluded.type,
                 kv=excluded.kv
            ''')
        self.c_add_record = self.db.add_statement(
            '''INSERT INTO RECORD VALUES (?, ?, ?, ?)
               ON CONFLICT DO UPDATE SET
                 name=excluded.name,
                 email=excluded.email
            ''')
        self.c_delete_issue = self.db.add_statement(
            'DELETE FROM ISSUES WHERE iid = ?')
        self.c_delete_record = self.db.add_statement(
            'DELETE FROM RECORD WHERE rid = ?')
        self.c_add_vote = self.db.add_statement(
            'INSERT INTO VOTES VALUES (NULL, ?, ?, ?, ?)')

        # Cursors for running queries.
        self.q_metadata = self.db.add_query('metadata',
            'SELECT * FROM METADATA')
        self.q_issues = self.db.add_query('issues',
            'SELECT * FROM ISSUES ORDER BY iid')
        self.q_record = self.db.add_query('record',
            'SELECT * FROM RECORD ORDER BY rid')
        self.q_get_issue = self.db.add_query('issues',
            'SELECT * FROM ISSUES WHERE iid = ?')
        self.q_get_record = self.db.add_query('record',
            'SELECT * FROM RECORD WHERE rid = ?')

    def open(self):

        # Double-check the Election is in the editing state.
        assert self.is_editable()

        # Add salts first. If this is gonna fail, then make sure it
        # happens before we move to the "opened" state.
        self.add_salts()

        edata = self.gather_election_data()
        print('EDATA:', edata)
        salt = crypto.gen_salt()
        opened_key = crypto.gen_opened_key(edata, salt)

        print('SALT:', salt)
        print('KEY:', opened_key)
        self.c_open.perform((salt, opened_key))

    def gather_election_data(self):
        "Gather a definition of this election for keying and anti-tamper."

        # NOTE: separators and other zero-entropy constant chars are
        # not included when assembling the data for hashing. This data
        # is not intended for human consumption, anyways.

        # NOTE: all assembly of rows must use a repeatable ordering.

        md = self.q_metadata.first_row()
        mdata = md.eid + md.title

        self.q_issues.perform()
        # Use an f-string to render "None" if a column is NULL.
        idata = ''.join(f'{r.iid}{r.title}{r.description}{r.type}{r.kv}'
                        for r in self.q_issues.fetchall())

        self.q_record.perform()
        rdata = ''.join(r.rid + r.email
                        for r in self.q_record.fetchall())

        return (mdata + idata + rdata).encode()

    def close(self):
        "Close an election."

        # The Election should be open.
        assert self.is_open()

        # Simple tweak of the metadata to close the Election.
        self.c_close.perform()

    def add_salts(self):
        "Set the SALT column in the ISSUES and RECORD tables."

        # The Election should be editable.
        assert self.is_editable()

        cur = self.db.conn.cursor()

        def for_table(table, mod_cursor):
            "Use MOD_CURSOR to salt each row of TABLE."

            # Fetch all ROWID values now, to avoid two cursors
            # attempting to work on TABLE at the same time.
            cur.execute(f'SELECT _ROWID_ FROM {table}')
            ids = cur.fetchall()

            # Now, add a salt to every row.
            for r in ids:
                salt = crypto.gen_salt()
                print('ROW will use:', table, r[0], salt)
                mod_cursor.perform((salt, r[0]))

        for_table('issues', self.c_salt_issue)
        for_table('record', self.c_salt_record)

    def get_issue(self, iid):
        "Return TITLE, DESCRIPTION, TYPE, and KV for issue IID."

        # NEVER return issue.salt
        issue = self.q_get_issue.first_row((iid,))
        return issue.title, issue.description, issue.type, issue.kv

    def add_issue(self, iid, title, description, type, kv):
        "Add or update an issue designated by IID."
        assert self.is_editable()

        # If we ADD, then SALT will be NULL. If we UPDATE, then it will not
        # be touched (it should be NULL).
        self.c_add_issue.perform((iid, title, description, type, kv, None))

    def delete_issue(self, iid):
        "Delete the Issue designated by IID."
        assert self.is_editable()

        self.c_delete_issue.perform((iid,))

    def get_participant(self, rid):
        "Return NAME, EMAIL for Participant on record RID."

        # NEVER return record.salt
        record = self.q_get_record.first_row((rid,))
        return record.name, record.email

    def add_participant(self, rid, name, email):
        "Add or update a Participant (voter of record) designated by RID."
        assert self.is_editable()

        # If we ADD, then SALT will be NULL. If we UPDATE, then it will not
        # be touched (it should be NULL).
        self.c_add_record.perform((rid, name, email, None))

    def delete_participant(self, rid):
        "Delete the Participant designated by RID."
        assert self.is_editable()

        self.c_delete_record.perform((rid,))

    def add_vote(self, rid, iid, votestring):
        "Add VOTESTRING as the (latest) vote by RID for IID."

        md = self.q_metadata.first_row()
        record = self.q_get_record.first_row((rid,))
        issue = self.q_get_issue.first_row((iid,))

        voter_token = crypto.gen_token(md.opened_key, rid, record.salt)
        #print('VOTER:', rid, record.salt, voter_token)
        issue_token = crypto.gen_token(md.opened_key, iid, issue.salt)
        #print('ISSUE:', iid, issue.salt, issue_token)

        salt, token = crypto.create_vote(voter_token, issue_token, votestring)
        #print('SALT:', salt)
        #print('TOKEN:', token)
        self.c_add_vote.perform((voter_token, issue_token, salt, token))

    def is_tampered(self):

        # The Election should be open.
        assert self.is_open()

        md = self.q_metadata.first_row()

        # Compute an opened_key based on the current data.
        edata = self.gather_election_data()
        opened_key = crypto.gen_opened_key(edata, md.salt)

        print('EDATA:', edata)
        print('SALT:', md.salt)
        print('KEY:', opened_key)

        # The computed key should be unchanged.
        return opened_key != md.opened_key

    def is_editable(self):
        "Can this Election be edited?"
        md = self.q_metadata.first_row()
        return md.salt is None and md.opened_key is None

    def is_open(self):
        "Is this Election open for voting?"
        md = self.q_metadata.first_row()
        return (md.salt is not None
                and md.opened_key is not None
                and md.closed in (None, 0))

    def is_closed(self):
        "Has this Election been closed?"
        md = self.q_metadata.first_row()
        return (md.salt is None
                and md.opened_key is None
                and md.closed == 1)

def new_eid():
    "Create a new ElectionID."

    # Use 4 bytes of a salt, for 32 bits.
    b = crypto.gen_salt()

    # Format into 8 hex characters.
    return f'{b[0]:02x}{b[1]:02x}{b[2]:02x}{b[3]:02x}'
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import re

from .utils import Utils
from .printer import Printer
from .objects import Model, ModelAttribute, ModelAPI, ModelOperation

USER = 'User'

IGNORED_ATTRIBUTES = ["_fetchers"]

ATTRIBUTE_MAPPING = {
    'global': 'globalMetadata'
}

IGNORED_RESOURCES = ['PublicNetworkMacro', 'NetworkLayout', 'InfrastructureConfig']

RESOURCE_MAPPING = {
    'SubNetwork': 'Subnet',
    'SubNetworkTemplate': 'SubnetTemplate',
    'PortStatus': 'MonitoringPort',
    'EnterpriseNetworkMacro': 'EnterpriseNetwork',
    'RedundantGWGrp': 'RedundancyGroup',
    'Service': 'ApplicationService',
    'IPBinding': 'IPReservation',
    'VPNConnect': 'VPNConnection',
    'QosPrimitive': 'QOS',  # Not working because resource name is 'qos' instead of 'quoss'
    'EgressQosPrimitive': 'EgressQOSPolicy',
    'EgressACLTemplateEntry': 'EgressACLEntryTemplate',
    'IngressACLTemplateEntry': 'IngressACLEntryTemplate',
    'IngressAdvancedForwardingTemplate': 'IngressAdvFwdTemplate',
    'IngressAdvancedForwardingTemplateEntry': 'IngressAdvFwdEntryTemplate',
    'AutoDiscGateway': 'AutoDiscoveredGateway',
    'VirtualMachine': 'VM',
    'Vlan': 'VLAN',
    'VlanTemplate': 'VLANTemplate'
}

PACKAGE_MAPPING = {
    '/alarm': 'Alarms',
    '/appd': 'Application Designer',
    '/common': 'Metadata',
    '/eventlog': 'Event Logs',
    '/gateway': 'Gateway Management',
    '/infrastructure': 'Infrastructure Profiles',
    '/job': 'Jobs',
    '/licensemgmt': 'Licensing',
    '/network': 'Core Networking',
    '/nsg': 'Gateway Management',
    '/policy': 'Policies',
    '/policy/acl': 'Security Policies',
    '/policy/qos': 'Policies',
    '/stats': 'Statistics',
    '/sysmon': 'System Monitoring',
    '/systemconfig': 'System Configuration',
    '/usermgmt': 'User Management',
    '/vm': 'Virtual Machines',
    '/vport': 'Core Networking',
    '/certificate': 'Certificate',
    '/cms': 'Cloud Management System',
    '/ipsec': 'IP Sec',
    '/keyserver': 'Key Server',
    '/vmware': 'VMware',
}

class ModelsProcessor(object):
    """ Process all models information that will be send to the writer """

    @classmethod
    def process(cls, resources):
        """ Prepare all resources

            Args:
                resources: A list of all resources to manage

            Returns:
                a processed list of resources

        """
        models = dict()

        for resource_name, resource in resources.iteritems():

            model = Model()
            model.description = resource['model']['description']
            ModelsProcessor._process_package(model=model, package=resource['model']['package'])
            ModelsProcessor._process_name(model=model, name=resource['model']['entityName'], resource_name=resource['model']['resourceName'])
            ModelsProcessor._process_apis(model=model, apis=resource['apis'], resources=resources)
            ModelsProcessor._process_attributes(model=model, attributes=resource['model']['attributes'])

            models[model.remote_name] = model

        Printer.success('Processed succeed for %s objects' % len(models))

        return models

    @classmethod
    def _process_package(cls, model, package):
        """ Process package name

        """
        if package in PACKAGE_MAPPING:
            model.package = PACKAGE_MAPPING[package]
        else:
            model.package = package

    @classmethod
    def _process_name(cls, model, name, resource_name):
        """ Compute the name and plural name of from the swagger
            model

            Args:
                model: the model processed
                name: the name from swagger

        """
        if name in RESOURCE_MAPPING:
            model.name = RESOURCE_MAPPING[name]
        else:
            model.name = name

        model.instance_name = Utils.get_python_name(model.name)
        model.instance_name = Utils.get_python_name(model.name)
        model.plural_name = Utils.get_plural_name(model.name)
        model.instance_plural_name = Utils.get_python_name(model.plural_name)
        model.remote_name = Utils.get_singular_name(resource_name)
        model.resource_name = resource_name

    @classmethod
    def _process_apis(cls, model, apis, resources):
        """ Process apis for the given model

            Args:
                model: the model processed
                apis: the list of apis availble for the current swagger model
                relations: dict containing all relations between resources

        """

        for path, api in apis['children'].iteritems():

            if api['resourceName'] == model.resource_name:
                continue

            names = filter(bool, re.split('/\{id\}?/?', path))

            child_resource_name = names[-1]
            child_rest_name = Utils.get_singular_name(names[-1])

            if child_rest_name.startswith('all'):
                entity_name = resources[child_rest_name[3:]]['model']['entityName']
            else:
                entity_name = resources[child_rest_name]['model']['entityName']

            model_api = ModelAPI()
            model_api.path = path
            model_api.resource_name = child_resource_name
            model_api.remote_name = child_rest_name
            model_api.plural_name = Utils.get_plural_name(entity_name)
            model_api.instance_plural_name = Utils.get_python_name(model_api.plural_name)

            for operation in api['operations']:
                model_operation = ModelOperation()
                model_operation.method = operation['method']
                model_api.operations.append(model_operation)

            model.apis['children'][path] = model_api

        for path, api in apis['parents'].iteritems():

            model_api = ModelAPI()
            model_api.path = path
            model_api.resource_name = api['resourceName']
            model_api.remote_name = api['RESTName']

            for operation in api['operations']:
                model_operation = ModelOperation()
                model_operation.method = operation['method']
                model_api.operations.append(model_operation)

            model.apis['parents'][path] = model_api

        for path, api in apis['self'].iteritems():

            model_api = ModelAPI()
            model_api.path = path

            for operation in api['operations']:
                model_operation = ModelOperation()
                model_operation.method = operation['method']
                model_api.operations.append(model_operation)

            model.apis['self'][path] = model_api

    @classmethod
    def _process_attribute_local_name(cls, name):
        """ Change local name according to the remote name

            Args:
                name: the remote name

            Returns:
                A beautiful python name

        """
        if name in ATTRIBUTE_MAPPING:
            return Utils.get_python_name(ATTRIBUTE_MAPPING[name])

        return Utils.get_python_name(name)

    @classmethod
    def _process_attributes(cls, model, attributes):
        """ Removes ignored attributes and update attribute type / name

            Args:
                model: the processed model
                properties: the list of properties to process

        """

        for name, attr in attributes.iteritems():

            attribute = ModelAttribute()
            attribute.remote_name = name
            attribute.local_name = cls._process_attribute_local_name(name)

            attribute.description = attr['description']
            attribute.type = attr['type']
            attribute.required = attr['required']
            attribute.unique_items = attr['uniqueItems']
            attribute.filterable = attr['filterable']
            attribute.readonly = attr['readonly']
            attribute.orderable = attr['orderable']
            attribute.creation_only = attr['creationOnly']
            attribute.autogenerated = attr['autogenerated']
            attribute.format = attr['format']
            attribute.min_length = attr['minLength']
            attribute.max_length = attr['maxLength']
            attribute.min_value = attr['minValue']
            attribute.max_value = attr['maxValue']
            attribute.allowed_chars = attr['allowedChars']
            attribute.allowed_choices = attr['allowedChoices']
            attribute.default_order = attr['defaultOrder']

            attribute.local_type = Utils.get_python_type_name(type_name=attribute.type, attribute_name=name, object_name=model.name)

            if attribute.local_type == 'time':
                model.has_time_attribute = True

            if attribute.local_type:
                model.attributes.append(attribute)
            else:
                # Simply ignore attributes otherwise...
                # CS 02/06/2015
                # Ignoring attribute enterprise of object InfrastructurePortProfile
                # Ignoring attribute gateway of object InfrastructureGatewayProfile
                # Ignoring attribute enterprise of object InfrastructureGatewayProfile
                Printer.log("Deliberately ignoring attribute %s of object %s" % (attribute.remote_name, model.name))
<commit_msg>Fixed child_rest_name<commit_after># -*- coding: utf-8 -*-

import re

from .utils import Utils
from .printer import Printer
from .objects import Model, ModelAttribute, ModelAPI, ModelOperation

USER = 'User'

IGNORED_ATTRIBUTES = ["_fetchers"]

ATTRIBUTE_MAPPING = {
    'global': 'globalMetadata'
}

IGNORED_RESOURCES = ['PublicNetworkMacro', 'NetworkLayout', 'InfrastructureConfig']

RESOURCE_MAPPING = {
    'SubNetwork': 'Subnet',
    'SubNetworkTemplate': 'SubnetTemplate',
    'PortStatus': 'MonitoringPort',
    'EnterpriseNetworkMacro': 'EnterpriseNetwork',
    'RedundantGWGrp': 'RedundancyGroup',
    'Service': 'ApplicationService',
    'IPBinding': 'IPReservation',
    'VPNConnect': 'VPNConnection',
    'QosPrimitive': 'QOS',  # Not working because resource name is 'qos' instead of 'quoss'
    'EgressQosPrimitive': 'EgressQOSPolicy',
    'EgressACLTemplateEntry': 'EgressACLEntryTemplate',
    'IngressACLTemplateEntry': 'IngressACLEntryTemplate',
    'IngressAdvancedForwardingTemplate': 'IngressAdvFwdTemplate',
    'IngressAdvancedForwardingTemplateEntry': 'IngressAdvFwdEntryTemplate',
    'AutoDiscGateway': 'AutoDiscoveredGateway',
    'VirtualMachine': 'VM',
    'Vlan': 'VLAN',
    'VlanTemplate': 'VLANTemplate'
}

PACKAGE_MAPPING = {
    '/alarm': 'Alarms',
    '/appd': 'Application Designer',
    '/common': 'Metadata',
    '/eventlog': 'Event Logs',
    '/gateway': 'Gateway Management',
    '/infrastructure': 'Infrastructure Profiles',
    '/job': 'Jobs',
    '/licensemgmt': 'Licensing',
    '/network': 'Core Networking',
    '/nsg': 'Gateway Management',
    '/policy': 'Policies',
    '/policy/acl': 'Security Policies',
    '/policy/qos': 'Policies',
    '/stats': 'Statistics',
    '/sysmon': 'System Monitoring',
    '/systemconfig': 'System Configuration',
    '/usermgmt': 'User Management',
    '/vm': 'Virtual Machines',
    '/vport': 'Core Networking',
    '/certificate': 'Certificate',
    '/cms': 'Cloud Management System',
    '/ipsec': 'IP Sec',
    '/keyserver': 'Key Server',
    '/vmware': 'VMware',
}

class ModelsProcessor(object):
    """ Process all models information that will be send to the writer """

    @classmethod
    def process(cls, resources):
        """ Prepare all resources

            Args:
                resources: A list of all resources to manage

            Returns:
                a processed list of resources

        """
        models = dict()

        for resource_name, resource in resources.iteritems():

            model = Model()
            model.description = resource['model']['description']
            ModelsProcessor._process_package(model=model, package=resource['model']['package'])
            ModelsProcessor._process_name(model=model, name=resource['model']['entityName'], resource_name=resource['model']['resourceName'])
            ModelsProcessor._process_apis(model=model, apis=resource['apis'], resources=resources)
            ModelsProcessor._process_attributes(model=model, attributes=resource['model']['attributes'])

            models[model.remote_name] = model

        Printer.success('Processed succeed for %s objects' % len(models))

        return models

    @classmethod
    def _process_package(cls, model, package):
        """ Process package name

        """
        if package in PACKAGE_MAPPING:
            model.package = PACKAGE_MAPPING[package]
        else:
            model.package = package

    @classmethod
    def _process_name(cls, model, name, resource_name):
        """ Compute the name and plural name of from the swagger
            model

            Args:
                model: the model processed
                name: the name from swagger

        """
        if name in RESOURCE_MAPPING:
            model.name = RESOURCE_MAPPING[name]
        else:
            model.name = name

        model.instance_name = Utils.get_python_name(model.name)
        model.instance_name = Utils.get_python_name(model.name)
        model.plural_name = Utils.get_plural_name(model.name)
        model.instance_plural_name = Utils.get_python_name(model.plural_name)
        model.remote_name = Utils.get_singular_name(resource_name)
        model.resource_name = resource_name

    @classmethod
    def _process_apis(cls, model, apis, resources):
        """ Process apis for the given model

            Args:
                model: the model processed
                apis: the list of apis availble for the current swagger model
                relations: dict containing all relations between resources

        """

        for path, api in apis['children'].iteritems():

            if api['resourceName'] == model.resource_name:
                continue

            names = filter(bool, re.split('/\{id\}?/?', path))

            child_resource_name = names[-1]
            child_rest_name = Utils.get_singular_name(names[-1])

            model_api = ModelAPI()
            model_api.path = path
            model_api.resource_name = child_resource_name
            model_api.remote_name = child_rest_name

            if child_rest_name.startswith('all'):
                child_rest_name = child_rest_name[3:]

            if child_rest_name in resources:
                entity_name = resources[child_rest_name]['model']['entityName']
                model_api.plural_name = Utils.get_plural_name(entity_name)
                model_api.instance_plural_name = Utils.get_python_name(model_api.plural_name)

            for operation in api['operations']:
                model_operation = ModelOperation()
                model_operation.method = operation['method']
                model_api.operations.append(model_operation)

            model.apis['children'][path] = model_api

        for path, api in apis['parents'].iteritems():

            model_api = ModelAPI()
            model_api.path = path
            model_api.resource_name = api['resourceName']
            model_api.remote_name = api['RESTName']

            for operation in api['operations']:
                model_operation = ModelOperation()
                model_operation.method = operation['method']
                model_api.operations.append(model_operation)

            model.apis['parents'][path] = model_api

        for path, api in apis['self'].iteritems():

            model_api = ModelAPI()
            model_api.path = path

            for operation in api['operations']:
                model_operation = ModelOperation()
                model_operation.method = operation['method']
                model_api.operations.append(model_operation)

            model.apis['self'][path] = model_api

    @classmethod
    def _process_attribute_local_name(cls, name):
        """ Change local name according to the remote name

            Args:
                name: the remote name

            Returns:
                A beautiful python name

        """
        if name in ATTRIBUTE_MAPPING:
            return Utils.get_python_name(ATTRIBUTE_MAPPING[name])

        return Utils.get_python_name(name)

    @classmethod
    def _process_attributes(cls, model, attributes):
        """ Removes ignored attributes and update attribute type / name

            Args:
                model: the processed model
                properties: the list of properties to process

        """

        for name, attr in attributes.iteritems():

            attribute = ModelAttribute()
            attribute.remote_name = name
            attribute.local_name = cls._process_attribute_local_name(name)

            attribute.description = attr['description']
            attribute.type = attr['type']
            attribute.required = attr['required']
            attribute.unique_items = attr['uniqueItems']
            attribute.filterable = attr['filterable']
            attribute.readonly = attr['readonly']
            attribute.orderable = attr['orderable']
            attribute.creation_only = attr['creationOnly']
            attribute.autogenerated = attr['autogenerated']
            attribute.format = attr['format']
            attribute.min_length = attr['minLength']
            attribute.max_length = attr['maxLength']
            attribute.min_value = attr['minValue']
            attribute.max_value = attr['maxValue']
            attribute.allowed_chars = attr['allowedChars']
            attribute.allowed_choices = attr['allowedChoices']
            attribute.default_order = attr['defaultOrder']

            attribute.local_type = Utils.get_python_type_name(type_name=attribute.type, attribute_name=name, object_name=model.name)

            if attribute.local_type == 'time':
                model.has_time_attribute = True

            if attribute.local_type:
                model.attributes.append(attribute)
            else:
                # Simply ignore attributes otherwise...
                # CS 02/06/2015
                # Ignoring attribute enterprise of object InfrastructurePortProfile
                # Ignoring attribute gateway of object InfrastructureGatewayProfile
                # Ignoring attribute enterprise of object InfrastructureGatewayProfile
                Printer.log("Deliberately ignoring attribute %s of object %s" % (attribute.remote_name, model.name))
<|endoftext|>


<commit_before>
# standard
import cStringIO
import optparse
import os
import resource
import signal
import string
import sys
import time
import traceback
from collections import defaultdict

# local
import sink
from core import __version__

# vendor
import gevent, gevent.socket
socket = gevent.socket
# protect stats
from gevent.thread import allocate_lock as Lock
stats_lock = Lock()

# constants
INTERVAL = 10.0
PERCENT = 90.0
MAX_PACKET = 2048

DESCRIPTION = '''
A statsd service in Python + gevent.
'''
EPILOG = '''

'''

# table to remove invalid characters from keys
ALL_ASCII = set(chr(c) for c in range(256))
KEY_VALID = string.ascii_letters + string.digits + '_-.'
KEY_TABLE = string.maketrans(KEY_VALID + '/', KEY_VALID + '_')
KEY_DELETIONS = ''.join(ALL_ASCII.difference(KEY_VALID + '/'))

# error messages
E_BADADDR = 'invalid bind address specified %r'
E_NOSINKS = 'you must specify at least one stats sink'

class Stats(object):

    def __init__(self):
        self.timers = defaultdict(list)
        self.counts = defaultdict(float)
        self.gauges = defaultdict(float)
        self.percent = PERCENT
        self.interval = INTERVAL

def daemonize(umask=0027):
    if gevent.fork():
        os._exit(0)
    os.setsid()
    if gevent.fork():
        os._exit(0)
    os.umask(umask)
    fd_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
    if fd_limit == resource.RLIM_INFINITY:
        fd_limit = 1024
    for fd in xrange(0, fd_limit):
        try:
            os.close(fd)
        except:
            pass
    os.open(os.devnull, os.O_RDWR)
    os.dup2(0, 1)
    os.dup2(0, 2)
    gevent.reinit()

def parse_addr(text):
    "Parse a 1- to 3-part address spec."
    if text:
        parts = text.split(':')
        length = len(parts)
        if length== 3:
            return parts[0], parts[1], int(parts[2])
        elif length == 2:
            return None, parts[0], int(parts[1])
        elif length == 1:
            return None, '', int(parts[0])
    return None, None, None

class StatsDaemon(object):

    """
    A statsd service implementation in Python + gevent.
    """

    def __init__(self, bindaddr, sinkspecs, interval, percent, debug=0,
                 key_prefix=''):
        _, host, port = parse_addr(bindaddr)
        if port is None:
            self.exit(E_BADADDR % bindaddr)
        self._bindaddr = (host, port)

        # TODO: generalize to support more than one sink type.  currently
        # only the graphite backend is present, but we may want to write
        # stats to hbase, redis, etc. - ph

        # construct the sink and add hosts to it
        if not sinkspecs:
            self.exit(E_NOSINKS)
        self._sink = sink.GraphiteSink()
        errors = []
        for spec in sinkspecs:
            try:
                self._sink.add(spec)
            except ValueError, ex:
                errors.append(ex)
        if errors:
            for err in errors:
                self.error(str(err))
            self.exit('exiting.')

        self._percent = float(percent)
        self._interval = float(interval)
        self._debug = debug
        self._sock = None
        self._flush_task = None
        self._key_prefix = key_prefix

        self._reset_stats()

    def _reset_stats(self):
        with stats_lock:
            self._stats = Stats()
            self._stats.percent = self._percent
            self._stats.interval = self._interval

    def exit(self, msg, code=1):
        self.error(msg)
        sys.exit(code)

    def error(self, msg):
        sys.stderr.write(msg + '\n')

    def start(self):
        "Start the service"
        # register signals
        gevent.signal(signal.SIGINT, self._shutdown)

        # spawn the flush trigger
        def _flush_impl():
            while 1:
                gevent.sleep(self._stats.interval)

                # rotate stats
                stats = self._stats
                self._reset_stats()

                # send the stats to the sink which in turn broadcasts
                # the stats packet to one or more hosts.
                try:
                    self._sink.send(stats)
                except Exception, ex:
                    trace = traceback.format_tb(sys.exc_info()[-1])
                    self.error(''.join(trace))

        self._flush_task = gevent.spawn(_flush_impl)

        # start accepting connections
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
            socket.IPPROTO_UDP)
        self._sock.bind(self._bindaddr)
        while 1:
            try:
                self._process(*self._sock.recvfrom(MAX_PACKET))
            except Exception, ex:
                self.error(str(ex))

    def _shutdown(self):
        "Shutdown the server"
        self.exit("service exiting", code=0)

    def _process(self, data, _):
        "Process a single packet and update the internal tables."
        parts = data.split(':')
        if self._debug:
            self.error('packet: %r' % data)
        if not parts:
            return

        # interpret the packet and update stats
        stats = self._stats
        key = parts[0].translate(KEY_TABLE, KEY_DELETIONS)
        if self._key_prefix:
            key = '.'.join([self._key_prefix, key])
        for part in parts[1:]:
            srate = 1.0
            fields = part.split('|')
            length = len(fields)
            if length < 2:
                continue
            value = fields[0]
            stype = fields[1].strip()

            with stats_lock:
                # timer (milliseconds)
                if stype == 'ms':
                    stats.timers[key].append(float(value if value else 0))

                # counter with optional sample rate
                elif stype == 'c':
                    if length == 3 and fields[2].startswith('@'):
                        srate = float(fields[2][1:])
                    value = float(value if value else 1) * (1 / srate)
                    stats.counts[key] += value
                elif stype == 'g':
                    value = float(value if value else 1)
                    stats.gauges[key] = value

def main():
    opts = optparse.OptionParser(description=DESCRIPTION, version=__version__,
        add_help_option=False)
    opts.add_option('-b', '--bind', dest='bind_addr', default=':8125',
        help="bind [host]:port (host defaults to '')")
    opts.add_option('-s', '--sink', dest='sink', action='append', default=[],
        help="a graphite service to which stats are sent ([host]:port).")
    opts.add_option('-v', dest='verbose', action='count', default=0,
        help="increase verbosity (currently used for debugging)")
    opts.add_option('-f', '--flush', dest='interval', default=INTERVAL,
        help="flush interval, in seconds (default 10)")
    opts.add_option('-x', '--prefix', dest='key_prefix', default='',
        help="key prefix added to all keys (default None)")
    opts.add_option('-p', '--percent', dest='percent', default=PERCENT,
        help="percent threshold (default 90)")
    opts.add_option('-D', '--daemonize', dest='daemonize', action='store_true',
        help='daemonize the service')
    opts.add_option('-h', '--help', dest='usage', action='store_true')

    (options, args) = opts.parse_args()

    if options.usage:
        # TODO: write epilog. usage is manually output since optparse will
        # wrap the epilog and we want pre-formatted output. - ph
        print(opts.format_help())
        sys.exit()

    if options.daemonize:
        daemonize()

    sd = StatsDaemon(options.bind_addr, options.sink, options.interval,
                     options.percent, options.verbose, options.key_prefix)
    sd.start()

if __name__ == '__main__':
    main()
<commit_msg>split received data on newlines<commit_after>
# standard
import cStringIO
import optparse
import os
import resource
import signal
import string
import sys
import time
import traceback
from collections import defaultdict

# local
import sink
from core import __version__

# vendor
import gevent, gevent.socket
socket = gevent.socket
# protect stats
from gevent.thread import allocate_lock as Lock
stats_lock = Lock()

# constants
INTERVAL = 10.0
PERCENT = 90.0
MAX_PACKET = 2048

DESCRIPTION = '''
A statsd service in Python + gevent.
'''
EPILOG = '''

'''

# table to remove invalid characters from keys
ALL_ASCII = set(chr(c) for c in range(256))
KEY_VALID = string.ascii_letters + string.digits + '_-.'
KEY_TABLE = string.maketrans(KEY_VALID + '/', KEY_VALID + '_')
KEY_DELETIONS = ''.join(ALL_ASCII.difference(KEY_VALID + '/'))

# error messages
E_BADADDR = 'invalid bind address specified %r'
E_NOSINKS = 'you must specify at least one stats sink'

class Stats(object):

    def __init__(self):
        self.timers = defaultdict(list)
        self.counts = defaultdict(float)
        self.gauges = defaultdict(float)
        self.percent = PERCENT
        self.interval = INTERVAL

def daemonize(umask=0027):
    if gevent.fork():
        os._exit(0)
    os.setsid()
    if gevent.fork():
        os._exit(0)
    os.umask(umask)
    fd_limit = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
    if fd_limit == resource.RLIM_INFINITY:
        fd_limit = 1024
    for fd in xrange(0, fd_limit):
        try:
            os.close(fd)
        except:
            pass
    os.open(os.devnull, os.O_RDWR)
    os.dup2(0, 1)
    os.dup2(0, 2)
    gevent.reinit()

def parse_addr(text):
    "Parse a 1- to 3-part address spec."
    if text:
        parts = text.split(':')
        length = len(parts)
        if length== 3:
            return parts[0], parts[1], int(parts[2])
        elif length == 2:
            return None, parts[0], int(parts[1])
        elif length == 1:
            return None, '', int(parts[0])
    return None, None, None

class StatsDaemon(object):

    """
    A statsd service implementation in Python + gevent.
    """

    def __init__(self, bindaddr, sinkspecs, interval, percent, debug=0,
                 key_prefix=''):
        _, host, port = parse_addr(bindaddr)
        if port is None:
            self.exit(E_BADADDR % bindaddr)
        self._bindaddr = (host, port)

        # TODO: generalize to support more than one sink type.  currently
        # only the graphite backend is present, but we may want to write
        # stats to hbase, redis, etc. - ph

        # construct the sink and add hosts to it
        if not sinkspecs:
            self.exit(E_NOSINKS)
        self._sink = sink.GraphiteSink()
        errors = []
        for spec in sinkspecs:
            try:
                self._sink.add(spec)
            except ValueError, ex:
                errors.append(ex)
        if errors:
            for err in errors:
                self.error(str(err))
            self.exit('exiting.')

        self._percent = float(percent)
        self._interval = float(interval)
        self._debug = debug
        self._sock = None
        self._flush_task = None
        self._key_prefix = key_prefix

        self._reset_stats()

    def _reset_stats(self):
        with stats_lock:
            self._stats = Stats()
            self._stats.percent = self._percent
            self._stats.interval = self._interval

    def exit(self, msg, code=1):
        self.error(msg)
        sys.exit(code)

    def error(self, msg):
        sys.stderr.write(msg + '\n')

    def start(self):
        "Start the service"
        # register signals
        gevent.signal(signal.SIGINT, self._shutdown)

        # spawn the flush trigger
        def _flush_impl():
            while 1:
                gevent.sleep(self._stats.interval)

                # rotate stats
                stats = self._stats
                self._reset_stats()

                # send the stats to the sink which in turn broadcasts
                # the stats packet to one or more hosts.
                try:
                    self._sink.send(stats)
                except Exception, ex:
                    trace = traceback.format_tb(sys.exc_info()[-1])
                    self.error(''.join(trace))

        self._flush_task = gevent.spawn(_flush_impl)

        # start accepting connections
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
            socket.IPPROTO_UDP)
        self._sock.bind(self._bindaddr)
        while 1:
            try:
                data, _ = self._sock.recvfrom(MAX_PACKET)
                for p in data.split('\n'):
                    if p:
                        self._process(p)
            except Exception, ex:
                self.error(str(ex))

    def _shutdown(self):
        "Shutdown the server"
        self.exit("service exiting", code=0)

    def _process(self, data):
        "Process a single packet and update the internal tables."
        parts = data.split(':')
        if self._debug:
            self.error('packet: %r' % data)
        if not parts:
            return

        # interpret the packet and update stats
        stats = self._stats
        key = parts[0].translate(KEY_TABLE, KEY_DELETIONS)
        if self._key_prefix:
            key = '.'.join([self._key_prefix, key])
        for part in parts[1:]:
            srate = 1.0
            fields = part.split('|')
            length = len(fields)
            if length < 2:
                continue
            value = fields[0]
            stype = fields[1].strip()

            with stats_lock:
                # timer (milliseconds)
                if stype == 'ms':
                    stats.timers[key].append(float(value if value else 0))

                # counter with optional sample rate
                elif stype == 'c':
                    if length == 3 and fields[2].startswith('@'):
                        srate = float(fields[2][1:])
                    value = float(value if value else 1) * (1 / srate)
                    stats.counts[key] += value
                elif stype == 'g':
                    value = float(value if value else 1)
                    stats.gauges[key] = value

def main():
    opts = optparse.OptionParser(description=DESCRIPTION, version=__version__,
        add_help_option=False)
    opts.add_option('-b', '--bind', dest='bind_addr', default=':8125',
        help="bind [host]:port (host defaults to '')")
    opts.add_option('-s', '--sink', dest='sink', action='append', default=[],
        help="a graphite service to which stats are sent ([host]:port).")
    opts.add_option('-v', dest='verbose', action='count', default=0,
        help="increase verbosity (currently used for debugging)")
    opts.add_option('-f', '--flush', dest='interval', default=INTERVAL,
        help="flush interval, in seconds (default 10)")
    opts.add_option('-x', '--prefix', dest='key_prefix', default='',
        help="key prefix added to all keys (default None)")
    opts.add_option('-p', '--percent', dest='percent', default=PERCENT,
        help="percent threshold (default 90)")
    opts.add_option('-D', '--daemonize', dest='daemonize', action='store_true',
        help='daemonize the service')
    opts.add_option('-h', '--help', dest='usage', action='store_true')

    (options, args) = opts.parse_args()

    if options.usage:
        # TODO: write epilog. usage is manually output since optparse will
        # wrap the epilog and we want pre-formatted output. - ph
        print(opts.format_help())
        sys.exit()

    if options.daemonize:
        daemonize()

    sd = StatsDaemon(options.bind_addr, options.sink, options.interval,
                     options.percent, options.verbose, options.key_prefix)
    sd.start()

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>"""Contains the BallController class which manages and tracks all the balls in
a pinball machine."""

import logging

from mpf.core.delays import DelayManager
from mpf.core.utility_functions import Util

class BallController(object):
    """Base class for the Ball Controller which is used to keep track of all
    the balls in a pinball machine.

    Parameters
    ----------

    machine : :class:`MachineController`
        A reference to the instance of the MachineController object.

    """
    def __init__(self, machine):
        self.machine = machine
        self.log = logging.getLogger("BallController")
        self.log.debug("Loading the BallController")
        self.delay = DelayManager(self.machine.delayRegistry)

        self.game = None

        self.num_balls_known = -999

        self.num_balls_missing = 0
        # Balls lost and/or not installed.

        # register for events
        self.machine.events.add_handler('request_to_start_game',
                                        self.request_to_start_game)
        self.machine.events.add_handler('machine_reset_phase_2',
                                        self._initialize)
        self.machine.events.add_handler('init_phase_2',
                                        self._init2)

    def _init2(self):
        # register a handler for all switches
        for device in self.machine.ball_devices:
            if 'ball_switches' not in device.config:
                continue
            for switch in device.config['ball_switches']:
                self.machine.switch_controller.add_switch_handler(switch.name,
                                                                  self._update_num_balls_known,
                                                                  ms=device.config['entrance_count_delay'])
                self.machine.switch_controller.add_switch_handler(switch.name,
                                                                  self._correct_playfield_count,
                                                                  ms=device.config['entrance_count_delay'],
                                                                  state=1)
                self.machine.switch_controller.add_switch_handler(switch.name,
                                                                  self._correct_playfield_count,
                                                                  ms=device.config['exit_count_delay'],
                                                                  state=0)

        for playfield in self.machine.playfields:
            self.machine.events.add_handler('{}_ball_count_change'.format(playfield.name),
                                            self._correct_playfield_count)

        # run initial count
        self._update_num_balls_known()

    def _get_loose_balls(self):
        return self.num_balls_known - self._count_stable_balls()

    def _count_stable_balls(self):
        self.log.debug("Counting Balls")
        balls = 0

        for device in self.machine.ball_devices:
            if device.is_playfield():
                continue

            if not device.is_ball_count_stable():
                raise ValueError("devices not stable")

            # generally we do not count ball devices without switches
            if 'ball_switches' not in device.config:
                continue
            # special handling for troughs (needed for gottlieb)
            elif not device.config['ball_switches'] and 'trough' in device.tags:
                balls += device.balls
            else:
                for switch in device.config['ball_switches']:
                    if self.machine.switch_controller.is_active(
                            switch.name, ms=device.config['entrance_count_delay']):
                        balls += 1
                    elif self.machine.switch_controller.is_inactive(
                            switch.name, ms=device.config['exit_count_delay']):
                        continue
                    else:
                        raise ValueError("switches not stable")

        return balls

    def _correct_playfield_count(self, **kwargs):
        del kwargs
        self.delay.reset(ms=1, callback=self._correct_playfield_count2, name="correct_playfield")

    def _correct_playfield_count2(self):
        try:
            loose_balls = self._get_loose_balls()
        except ValueError:
            self.delay.reset(ms=10000, callback=self._correct_playfield_count2, name="correct_playfield")
            return

        balls_on_pfs = 0

        for playfield in self.machine.playfields:
            balls_on_pfs += playfield.balls

        # fix too much balls
        while balls_on_pfs > loose_balls:
            for playfield in self.machine.playfields:
                self.log.warning("Corecting balls on pf from %s to %s",
                                 playfield.balls, playfield.balls - 1)
                if playfield.balls > 0:
                    if playfield.unexpected_balls > 0:
                        playfield.unexpected_balls -= 1
                    playfield.balls -= 1
                    balls_on_pfs -= 1
                    break

        for playfield in self.machine.playfields:
            if playfield.balls != playfield.available_balls:
                self.log.warning("Corecting available_balls %s to %s",
                                 playfield.available_balls, playfield.balls)
                playfield.available_balls = playfield.balls

    def _update_num_balls_known(self):
        balls = self._count_balls()
        if balls > self.num_balls_known:
            self.log.debug("Found new balls. Setting known balls to %s", balls)
            self.delay.add(1, self._handle_new_balls, new_balls=balls - self.num_balls_known)
            self.num_balls_known = balls

    def _handle_new_balls(self, new_balls):
        for dummy_i in range(new_balls):
            for playfield in self.machine.playfields:
                if playfield.unexpected_balls > 0:
                    playfield.unexpected_balls -= 1
                    break

    def _count_balls(self):
        self.log.debug("Counting Balls")
        balls = 0

        for device in self.machine.ball_devices:
            # generally we do not count ball devices without switches
            if 'ball_switches' not in device.config:
                continue
            # special handling for troughs (needed for gottlieb)
            elif not device.config['ball_switches'] and 'trough' in device.tags:
                balls += device.balls
            else:
                for switch in device.config['ball_switches']:
                    if self.machine.switch_controller.is_active(
                            switch.name, ms=device.config['entrance_count_delay']):
                        balls += 1

        return balls

    def _initialize(self):

        # If there are no ball devices, then the ball controller has no work to
        # do and will create errors, so we just abort.
        if not hasattr(self.machine, 'ball_devices'):
            return

        for device in self.machine.ball_devices:
            if 'drain' in device.tags:  # device is used to drain balls from pf
                self.machine.events.add_handler('balldevice_' + device.name +
                                                '_ball_enter',
                                                self._ball_drained_handler)

    def request_to_start_game(self):
        """Method registered for the *request_to_start_game* event.

        Checks to make sure that the balls are in all the right places and
        returns. If too many balls are missing (based on the config files 'Min
        Balls' setting), it will return False to reject the game start request.
        """
        balls = self._count_balls()
        self.log.debug("Received request to start game.")
        self.log.debug("Balls contained: %s, Min balls needed: %s",
                       balls,
                       self.machine.config['machine']['min_balls'])
        if balls < self.machine.config['machine']['min_balls']:
            self.log.warning("BallController denies game start. Not enough "
                             "balls")
            return False

        if self.machine.config['game']['allow_start_with_ball_in_drain']:
            allowed_positions = ['home', 'trough', 'drain']
        else:
            allowed_positions = ['home', 'trough']

        if self.machine.config['game']['allow_start_with_loose_balls']:
            return

        elif not self.are_balls_collected(allowed_positions):
            self.collect_balls('home')
            self.log.warning("BallController denies game start. Balls are not "
                             "in their home positions.")
            return False

    def are_balls_collected(self, target):
        """Checks to see if all the balls are contained in devices tagged with
        the parameter that was passed.

        Note if you pass a target that's not used in any ball devices, this
        method will return True. (Because you're asking if all balls are
        nowhere, and they always are. :)

        Args:
            target: String or list of strings of the tags you'd like to
                collect the balls to. Default of None will be replaced with
                'home' and 'trough'.

        """

        self.log.debug("Checking to see if all the balls are in devices tagged"
                       " with '%s'", target)

        if isinstance(target, str):
            target = Util.string_to_list(target)

        count = 0
        devices = set()

        for tag in target:
            for device in self.machine.ball_devices.items_tagged(tag):
                devices.add(device)

        if len(devices) == 0:
            # didn't find any devices matching that tag, so we return True
            return True

        for device in devices:
            count += device.get_status('balls')
            self.log.debug('Found %s ball(s) in %s. Found %s total',
                           device.get_status('balls'), device.name, count)

        if count == self.machine.ball_controller.num_balls_known:
            self.log.debug("Yes, all balls are collected")
            return True
        else:
            self.log.debug("No, all balls are not collected. Balls Counted: %s. "
                           "Total balls known: %s", count,
                           self.machine.ball_controller.num_balls_known)
            return False

    def collect_balls(self, target='home, trough'):
        """Used to ensure that all balls are in contained in ball devices with
        the tag or list of tags you pass.

        Typically this would be used after a game ends, or when the machine is
        reset or first starts up, to ensure that all balls are in devices
        tagged with 'home' and/or 'trough'.

        Args:
            target: A string of the tag name or a list of tags names of the
                ball devices you want all the balls to end up in. Default is
                ['home', 'trough'].

        """
        # I'm embarrassed at how ugly this code is. But meh, it works...

        tag_list = Util.string_to_list(target)

        self.log.debug("Collecting all balls to devices with tags '%s'",
                       tag_list)

        target_devices = set()
        source_devices = set()
        balls_to_collect = False

        for tag in tag_list:
            for device in self.machine.ball_devices.items_tagged(tag):
                target_devices.add(device)

        for device in self.machine.ball_devices:
            if device not in target_devices:
                if device.available_balls:
                    source_devices.add(device)
                    balls_to_collect = True

        self.log.debug("Ejecting all balls from: %s", source_devices)

        if balls_to_collect:
            self.machine.events.post('collecting_balls')
            '''event: collecting_balls

            desc: Posted by the ball controller when it starts the collecting
                balls process.

            '''

            for device in target_devices:
                self.machine.events.replace_handler(
                    'balldevice_{}_ball_enter'.format(device.name),
                    self._collecting_balls_entered_callback,
                    target=target)

            for device in source_devices:
                if not device.is_playfield():
                    device.eject_all()
        else:
            self.log.debug("All balls are collected")
            self._collecting_balls_complete()

    def _collecting_balls_entered_callback(self, target, new_balls, unclaimed_balls, **kwargs):
        del kwargs
        del new_balls
        if self.are_balls_collected(target=target):
            self._collecting_balls_complete()

        return {'unclaimed_balls': unclaimed_balls}

    def _collecting_balls_complete(self):
        self.machine.events.remove_handler(self._collecting_balls_complete)
        self.machine.events.post('collecting_balls_complete')
        '''event: collecting_balls_complete

        desc: Posted by the ball controller when it has finished the collecting
            balls process.

        '''

    def _ball_drained_handler(self, new_balls, unclaimed_balls, device, **kwargs):
        del kwargs
        del new_balls
        self.machine.events.post_relay('ball_drain',
                                       callback=self._process_ball_drained,
                                       device=device,
                                       balls=unclaimed_balls)

        # What happens if the ball enters the trough but the ball_add_live
        # event hasn't confirmed its eject? todo

    def _process_ball_drained(self, balls=None, ev_result=None, **kwargs):
        # We don't need to do anything here because other modules (ball save,
        # the game, etc. should jump in and do whatever they need to do when a
        # ball is drained.
        pass
<commit_msg>fix potential endless loop<commit_after>"""Contains the BallController class which manages and tracks all the balls in
a pinball machine."""

import logging

from mpf.core.delays import DelayManager
from mpf.core.utility_functions import Util

class BallController(object):
    """Base class for the Ball Controller which is used to keep track of all
    the balls in a pinball machine.

    Parameters
    ----------

    machine : :class:`MachineController`
        A reference to the instance of the MachineController object.

    """
    def __init__(self, machine):
        self.machine = machine
        self.log = logging.getLogger("BallController")
        self.log.debug("Loading the BallController")
        self.delay = DelayManager(self.machine.delayRegistry)

        self.game = None

        self.num_balls_known = -999

        self.num_balls_missing = 0
        # Balls lost and/or not installed.

        # register for events
        self.machine.events.add_handler('request_to_start_game',
                                        self.request_to_start_game)
        self.machine.events.add_handler('machine_reset_phase_2',
                                        self._initialize)
        self.machine.events.add_handler('init_phase_2',
                                        self._init2)

    def _init2(self):
        # register a handler for all switches
        for device in self.machine.ball_devices:
            if 'ball_switches' not in device.config:
                continue
            for switch in device.config['ball_switches']:
                self.machine.switch_controller.add_switch_handler(switch.name,
                                                                  self._update_num_balls_known,
                                                                  ms=device.config['entrance_count_delay'])
                self.machine.switch_controller.add_switch_handler(switch.name,
                                                                  self._correct_playfield_count,
                                                                  ms=device.config['entrance_count_delay'],
                                                                  state=1)
                self.machine.switch_controller.add_switch_handler(switch.name,
                                                                  self._correct_playfield_count,
                                                                  ms=device.config['exit_count_delay'],
                                                                  state=0)

        for playfield in self.machine.playfields:
            self.machine.events.add_handler('{}_ball_count_change'.format(playfield.name),
                                            self._correct_playfield_count)

        # run initial count
        self._update_num_balls_known()

    def _get_loose_balls(self):
        return self.num_balls_known - self._count_stable_balls()

    def _count_stable_balls(self):
        self.log.debug("Counting Balls")
        balls = 0

        for device in self.machine.ball_devices:
            if device.is_playfield():
                continue

            if not device.is_ball_count_stable():
                raise ValueError("devices not stable")

            # generally we do not count ball devices without switches
            if 'ball_switches' not in device.config:
                continue
            # special handling for troughs (needed for gottlieb)
            elif not device.config['ball_switches'] and 'trough' in device.tags:
                balls += device.balls
            else:
                for switch in device.config['ball_switches']:
                    if self.machine.switch_controller.is_active(
                            switch.name, ms=device.config['entrance_count_delay']):
                        balls += 1
                    elif self.machine.switch_controller.is_inactive(
                            switch.name, ms=device.config['exit_count_delay']):
                        continue
                    else:
                        raise ValueError("switches not stable")

        return balls

    def _correct_playfield_count(self, **kwargs):
        del kwargs
        self.delay.reset(ms=1, callback=self._correct_playfield_count2, name="correct_playfield")

    def _correct_playfield_count2(self):
        try:
            loose_balls = self._get_loose_balls()
        except ValueError:
            self.delay.reset(ms=10000, callback=self._correct_playfield_count2, name="correct_playfield")
            return

        balls_on_pfs = 0

        for playfield in self.machine.playfields:
            balls_on_pfs += playfield.balls

        # fix too much balls
        if balls_on_pfs > loose_balls:
            for dummy_i in range(balls_on_pfs - loose_balls):
                for playfield in self.machine.playfields:
                    self.log.warning("Corecting balls on pf from %s to %s",
                                     playfield.balls, playfield.balls - 1)
                    if playfield.balls > 0:
                        if playfield.unexpected_balls > 0:
                            playfield.unexpected_balls -= 1
                        playfield.balls -= 1
                        balls_on_pfs -= 1
                        break

        for playfield in self.machine.playfields:
            if playfield.balls != playfield.available_balls:
                self.log.warning("Corecting available_balls %s to %s",
                                 playfield.available_balls, playfield.balls)
                playfield.available_balls = playfield.balls

    def _update_num_balls_known(self):
        balls = self._count_balls()
        if balls > self.num_balls_known:
            self.log.debug("Found new balls. Setting known balls to %s", balls)
            self.delay.add(1, self._handle_new_balls, new_balls=balls - self.num_balls_known)
            self.num_balls_known = balls

    def _handle_new_balls(self, new_balls):
        for dummy_i in range(new_balls):
            for playfield in self.machine.playfields:
                if playfield.unexpected_balls > 0:
                    playfield.unexpected_balls -= 1
                    break

    def _count_balls(self):
        self.log.debug("Counting Balls")
        balls = 0

        for device in self.machine.ball_devices:
            # generally we do not count ball devices without switches
            if 'ball_switches' not in device.config:
                continue
            # special handling for troughs (needed for gottlieb)
            elif not device.config['ball_switches'] and 'trough' in device.tags:
                balls += device.balls
            else:
                for switch in device.config['ball_switches']:
                    if self.machine.switch_controller.is_active(
                            switch.name, ms=device.config['entrance_count_delay']):
                        balls += 1

        return balls

    def _initialize(self):

        # If there are no ball devices, then the ball controller has no work to
        # do and will create errors, so we just abort.
        if not hasattr(self.machine, 'ball_devices'):
            return

        for device in self.machine.ball_devices:
            if 'drain' in device.tags:  # device is used to drain balls from pf
                self.machine.events.add_handler('balldevice_' + device.name +
                                                '_ball_enter',
                                                self._ball_drained_handler)

    def request_to_start_game(self):
        """Method registered for the *request_to_start_game* event.

        Checks to make sure that the balls are in all the right places and
        returns. If too many balls are missing (based on the config files 'Min
        Balls' setting), it will return False to reject the game start request.
        """
        balls = self._count_balls()
        self.log.debug("Received request to start game.")
        self.log.debug("Balls contained: %s, Min balls needed: %s",
                       balls,
                       self.machine.config['machine']['min_balls'])
        if balls < self.machine.config['machine']['min_balls']:
            self.log.warning("BallController denies game start. Not enough "
                             "balls")
            return False

        if self.machine.config['game']['allow_start_with_ball_in_drain']:
            allowed_positions = ['home', 'trough', 'drain']
        else:
            allowed_positions = ['home', 'trough']

        if self.machine.config['game']['allow_start_with_loose_balls']:
            return

        elif not self.are_balls_collected(allowed_positions):
            self.collect_balls('home')
            self.log.warning("BallController denies game start. Balls are not "
                             "in their home positions.")
            return False

    def are_balls_collected(self, target):
        """Checks to see if all the balls are contained in devices tagged with
        the parameter that was passed.

        Note if you pass a target that's not used in any ball devices, this
        method will return True. (Because you're asking if all balls are
        nowhere, and they always are. :)

        Args:
            target: String or list of strings of the tags you'd like to
                collect the balls to. Default of None will be replaced with
                'home' and 'trough'.

        """

        self.log.debug("Checking to see if all the balls are in devices tagged"
                       " with '%s'", target)

        if isinstance(target, str):
            target = Util.string_to_list(target)

        count = 0
        devices = set()

        for tag in target:
            for device in self.machine.ball_devices.items_tagged(tag):
                devices.add(device)

        if len(devices) == 0:
            # didn't find any devices matching that tag, so we return True
            return True

        for device in devices:
            count += device.get_status('balls')
            self.log.debug('Found %s ball(s) in %s. Found %s total',
                           device.get_status('balls'), device.name, count)

        if count == self.machine.ball_controller.num_balls_known:
            self.log.debug("Yes, all balls are collected")
            return True
        else:
            self.log.debug("No, all balls are not collected. Balls Counted: %s. "
                           "Total balls known: %s", count,
                           self.machine.ball_controller.num_balls_known)
            return False

    def collect_balls(self, target='home, trough'):
        """Used to ensure that all balls are in contained in ball devices with
        the tag or list of tags you pass.

        Typically this would be used after a game ends, or when the machine is
        reset or first starts up, to ensure that all balls are in devices
        tagged with 'home' and/or 'trough'.

        Args:
            target: A string of the tag name or a list of tags names of the
                ball devices you want all the balls to end up in. Default is
                ['home', 'trough'].

        """
        # I'm embarrassed at how ugly this code is. But meh, it works...

        tag_list = Util.string_to_list(target)

        self.log.debug("Collecting all balls to devices with tags '%s'",
                       tag_list)

        target_devices = set()
        source_devices = set()
        balls_to_collect = False

        for tag in tag_list:
            for device in self.machine.ball_devices.items_tagged(tag):
                target_devices.add(device)

        for device in self.machine.ball_devices:
            if device not in target_devices:
                if device.available_balls:
                    source_devices.add(device)
                    balls_to_collect = True

        self.log.debug("Ejecting all balls from: %s", source_devices)

        if balls_to_collect:
            self.machine.events.post('collecting_balls')
            '''event: collecting_balls

            desc: Posted by the ball controller when it starts the collecting
                balls process.

            '''

            for device in target_devices:
                self.machine.events.replace_handler(
                    'balldevice_{}_ball_enter'.format(device.name),
                    self._collecting_balls_entered_callback,
                    target=target)

            for device in source_devices:
                if not device.is_playfield():
                    device.eject_all()
        else:
            self.log.debug("All balls are collected")
            self._collecting_balls_complete()

    def _collecting_balls_entered_callback(self, target, new_balls, unclaimed_balls, **kwargs):
        del kwargs
        del new_balls
        if self.are_balls_collected(target=target):
            self._collecting_balls_complete()

        return {'unclaimed_balls': unclaimed_balls}

    def _collecting_balls_complete(self):
        self.machine.events.remove_handler(self._collecting_balls_complete)
        self.machine.events.post('collecting_balls_complete')
        '''event: collecting_balls_complete

        desc: Posted by the ball controller when it has finished the collecting
            balls process.

        '''

    def _ball_drained_handler(self, new_balls, unclaimed_balls, device, **kwargs):
        del kwargs
        del new_balls
        self.machine.events.post_relay('ball_drain',
                                       callback=self._process_ball_drained,
                                       device=device,
                                       balls=unclaimed_balls)

        # What happens if the ball enters the trough but the ball_add_live
        # event hasn't confirmed its eject? todo

    def _process_ball_drained(self, balls=None, ev_result=None, **kwargs):
        # We don't need to do anything here because other modules (ball save,
        # the game, etc. should jump in and do whatever they need to do when a
        # ball is drained.
        pass
<|endoftext|>


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

#  Copyright (c) 2014-present, Facebook, Inc.
#  All rights reserved.
#
#  This source code is licensed under both the Apache 2.0 license (found in the
#  LICENSE file in the root directory of this source tree) and the GPLv2 (found
#  in the COPYING file in the root directory of this source tree).
#  You may select, at your option, one of the above-listed licenses.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import re
import sys
import argparse
import platform
import subprocess

ORACLE_RELEASE = "/etc/oracle-release"
SYSTEM_RELEASE = "/etc/system-release"
LSB_RELEASE    = "/etc/lsb-release"
OS_RELEASE     = "/etc/os-release"
DEBIAN_VERSION = "/etc/debian_version"
GENTOO_RELEASE = "/etc/gentoo-release"
SUSE_RELEASE   = "/etc/SuSE-release"

def _platform():
    osType, _, _, _, _, _ = platform.uname()

    if osType == "Windows":
        return ("windows", "windows")
    elif osType == "Linux":
        if os.path.exists(ORACLE_RELEASE):
            return ("redhat", "oracle")

        if os.path.exists(SYSTEM_RELEASE):
            with open(SYSTEM_RELEASE, "r") as fd:
                fileContents = fd.read()

                if fileContents.find("CentOS") != -1:
                    return ("redhat", "centos")

                if fileContents.find("Scientific Linux") != -1:
                    return ("redhat", "scientific")

                if fileContents.find("Red Hat Enterprise") != -1:
                    return ("redhat", "rhel")

                if fileContents.find("Amazon Linux") != -1:
                    return ("redhat", "amazon")

                if fileContents.find("Fedora") != -1:
                    return ("redhat", "fedora")

        if os.path.exists(LSB_RELEASE):
            with open(LSB_RELEASE, "r") as fd:
                fileContents = fd.read()

                if fileContents.find("DISTRIB_ID=Ubuntu") != -1:
                    return ("debian", "ubuntu")

                if fileContents.find("DISTRIB_ID=Arch") != -1:
                    return ("arch", "arch")

                if fileContents.find("DISTRIB_ID=ManjaroLinux") != -1:
                    return ("arch", "manjaro")

        if os.path.exists(OS_RELEASE):
            with open(OS_RELEASE, "r") as fd:
                fileContents = fd.read()

                if fileContents.find("ID=nixos") != -1:
                    return ("nixos", "nixos")

        if os.path.exists(DEBIAN_VERSION):
            return ("debian", "debian")

        if os.path.exists(GENTOO_RELEASE):
            return ("gentoo", "gentoo")

        if os.path.exists(SUSE_RELEASE):
            return ("suse", "suse")
    else:
        return (None, osType.lower())

def _distro(osType):
    def getRedhatDistroVersion(pattern):
        with open(SYSTEM_RELEASE, "r") as fd:
            contents = fd.read()

            result = re.findall(pattern, contents)
            if result and len(result) == 1:
                return result[0].replace("release ", osType)
        return None

    def commandOutput(cmd):
        try:
            output = subprocess.check_output(cmd)
            return output
        except subprocess.CalledProcessError:
            return None
        except OSError:
            return None
        except WindowsError:
            return None

    _, _, osVersion, _, _, _ = platform.uname()

    if osType == "oracle":
        result = getRedhatDistroVersion(r'release [5-7]')
        if result is not None:
            return result
    elif osType in ["centos", "scientific", "rhel"]:
        result = getRedhatDistroVersion(r'release [6-7]')
        if result is not None:
            return result
    elif osType == "amazon":
        result = getRedhatDistroVersion(r'release 20[12][0-9]\.[0-9][0-9]')
        if result is not None:
            return result
    elif osType == "ubuntu":
        with open(LSB_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'DISTRIB_CODENAME=(.*)', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "darwin":
        rawResult = commandOutput(["sw_vers", "-productVersion"])
        if rawResult is not None:
            results = re.findall(r'[0-9]+\.[0-9]+', rawResult)
            if len(results) == 1:
                return results[0]
    elif osType == "fedora":
        with open(SYSTEM_RELEASE, "r") as fd:
          contents = fd.read()
          results = contents.split()
          if len(results) > 2:
            return results[2]
    elif osType == "arch":
        with open("/etc/arch-release", "r") as fd:
          contents = fd.read()
          results = contents.split()
          if len(results) > 2:
            return results[2]
    elif osType == "manjaro":
        with open(LSB_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'DISTRIB_CODENAME=(.*)', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "debian":
        result = commandOutput(["lsb_release", "-cs"])
        if result is not None:
            return result
    elif osType == "freebsd":
        rawResult = commandOutput(["uname", "-r"])
        results = rawResult.split("-")
        if len(results) > 0:
          return results[0]
    elif osType == "gentoo":
        with open(GENTOO_RELEASE, "r") as fd:
          contents = fd.read()
          results = contents.split()
        if len(results) > 0:
          return results[len(results) -1]
    elif osType == "suse":
        with open(SUSE_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'VERSION = (.*)', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "nixos":
        with open(OS_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'VERSION_ID=\"(.*)\"', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "windows":
        return "windows%s" % osVersion

    return "unknown_version"

def platformAction():
    family, osType = _platform()
    print(osType)

def distroAction():
    family, osType = _platform()
    print(_distro(osType))

def familyAction():
    family, osType = _platform()
    if family:
        print(family)

def defaultAction():
    family, osType = _platform()
    distro = _distro(osType)
    print("%s;%s" % (osType, distro))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Platform detection script for osquery")
    parser.add_argument("--platform", action="store_true", help="Outputs the detected platform")
    parser.add_argument("--distro", action="store_true", help="Outputs the detected distribution")
    parser.add_argument("--family", action="store_true", help="Outputs the detected family")

    args = parser.parse_args()

    if args.platform and \
        not args.distro and \
        not args.family:
      platformAction()
    elif not args.platform and \
        args.distro and \
        not args.family:
      distroAction()
    elif not args.platform and \
        not args.distro and \
        args.family:
      familyAction()
    else:
      defaultAction()

    sys.exit(0)
<commit_msg>Prefer /etc/os-release to identify Arch Linux (#4759)<commit_after>#!/usr/bin/env python

#  Copyright (c) 2014-present, Facebook, Inc.
#  All rights reserved.
#
#  This source code is licensed under both the Apache 2.0 license (found in the
#  LICENSE file in the root directory of this source tree) and the GPLv2 (found
#  in the COPYING file in the root directory of this source tree).
#  You may select, at your option, one of the above-listed licenses.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import re
import sys
import argparse
import platform
import subprocess

ORACLE_RELEASE = "/etc/oracle-release"
SYSTEM_RELEASE = "/etc/system-release"
LSB_RELEASE    = "/etc/lsb-release"
OS_RELEASE     = "/etc/os-release"
DEBIAN_VERSION = "/etc/debian_version"
GENTOO_RELEASE = "/etc/gentoo-release"
SUSE_RELEASE   = "/etc/SuSE-release"

def _platform():
    osType, _, _, _, _, _ = platform.uname()

    if osType == "Windows":
        return ("windows", "windows")
    elif osType == "Linux":
        if os.path.exists(ORACLE_RELEASE):
            return ("redhat", "oracle")

        if os.path.exists(SYSTEM_RELEASE):
            with open(SYSTEM_RELEASE, "r") as fd:
                fileContents = fd.read()

                if fileContents.find("CentOS") != -1:
                    return ("redhat", "centos")

                if fileContents.find("Scientific Linux") != -1:
                    return ("redhat", "scientific")

                if fileContents.find("Red Hat Enterprise") != -1:
                    return ("redhat", "rhel")

                if fileContents.find("Amazon Linux") != -1:
                    return ("redhat", "amazon")

                if fileContents.find("Fedora") != -1:
                    return ("redhat", "fedora")

        if os.path.exists(LSB_RELEASE):
            with open(LSB_RELEASE, "r") as fd:
                fileContents = fd.read()

                if fileContents.find("DISTRIB_ID=Ubuntu") != -1:
                    return ("debian", "ubuntu")

                if fileContents.find("DISTRIB_ID=ManjaroLinux") != -1:
                    return ("arch", "manjaro")

        if os.path.exists(OS_RELEASE):
            with open(OS_RELEASE, "r") as fd:
                fileContents = fd.read()

                if fileContents.find("ID=arch") != -1:
                    return ("arch", "arch")

                if fileContents.find("ID=nixos") != -1:
                    return ("nixos", "nixos")

        if os.path.exists(DEBIAN_VERSION):
            return ("debian", "debian")

        if os.path.exists(GENTOO_RELEASE):
            return ("gentoo", "gentoo")

        if os.path.exists(SUSE_RELEASE):
            return ("suse", "suse")
    else:
        return (None, osType.lower())

def _distro(osType):
    def getRedhatDistroVersion(pattern):
        with open(SYSTEM_RELEASE, "r") as fd:
            contents = fd.read()

            result = re.findall(pattern, contents)
            if result and len(result) == 1:
                return result[0].replace("release ", osType)
        return None

    def commandOutput(cmd):
        try:
            output = subprocess.check_output(cmd)
            return output
        except subprocess.CalledProcessError:
            return None
        except OSError:
            return None
        except WindowsError:
            return None

    _, _, osVersion, _, _, _ = platform.uname()

    if osType == "oracle":
        result = getRedhatDistroVersion(r'release [5-7]')
        if result is not None:
            return result
    elif osType in ["centos", "scientific", "rhel"]:
        result = getRedhatDistroVersion(r'release [6-7]')
        if result is not None:
            return result
    elif osType == "amazon":
        result = getRedhatDistroVersion(r'release 20[12][0-9]\.[0-9][0-9]')
        if result is not None:
            return result
    elif osType == "ubuntu":
        with open(LSB_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'DISTRIB_CODENAME=(.*)', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "darwin":
        rawResult = commandOutput(["sw_vers", "-productVersion"])
        if rawResult is not None:
            results = re.findall(r'[0-9]+\.[0-9]+', rawResult)
            if len(results) == 1:
                return results[0]
    elif osType == "fedora":
        with open(SYSTEM_RELEASE, "r") as fd:
          contents = fd.read()
          results = contents.split()
          if len(results) > 2:
            return results[2]
    elif osType == "arch":
        with open("/etc/arch-release", "r") as fd:
          contents = fd.read()
          results = contents.split()
          if len(results) > 2:
            return results[2]
    elif osType == "manjaro":
        with open(LSB_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'DISTRIB_CODENAME=(.*)', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "debian":
        result = commandOutput(["lsb_release", "-cs"])
        if result is not None:
            return result
    elif osType == "freebsd":
        rawResult = commandOutput(["uname", "-r"])
        results = rawResult.split("-")
        if len(results) > 0:
          return results[0]
    elif osType == "gentoo":
        with open(GENTOO_RELEASE, "r") as fd:
          contents = fd.read()
          results = contents.split()
        if len(results) > 0:
          return results[len(results) -1]
    elif osType == "suse":
        with open(SUSE_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'VERSION = (.*)', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "nixos":
        with open(OS_RELEASE, "r") as fd:
            contents = fd.read()
            results = re.findall(r'VERSION_ID=\"(.*)\"', contents)
            if len(results) == 1:
                return results[0]
    elif osType == "windows":
        return "windows%s" % osVersion

    return "unknown_version"

def platformAction():
    family, osType = _platform()
    print(osType)

def distroAction():
    family, osType = _platform()
    print(_distro(osType))

def familyAction():
    family, osType = _platform()
    if family:
        print(family)

def defaultAction():
    family, osType = _platform()
    distro = _distro(osType)
    print("%s;%s" % (osType, distro))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Platform detection script for osquery")
    parser.add_argument("--platform", action="store_true", help="Outputs the detected platform")
    parser.add_argument("--distro", action="store_true", help="Outputs the detected distribution")
    parser.add_argument("--family", action="store_true", help="Outputs the detected family")

    args = parser.parse_args()

    if args.platform and \
        not args.distro and \
        not args.family:
      platformAction()
    elif not args.platform and \
        args.distro and \
        not args.family:
      distroAction()
    elif not args.platform and \
        not args.distro and \
        args.family:
      familyAction()
    else:
      defaultAction()

    sys.exit(0)
<|endoftext|>