<commit_before>class Dcim(object):

    def __init__(self, netbox_con):
        self.netbox_con = netbox_con

    def get_sites(self):
        """Returns the available sites"""
        return self.netbox_con.get('/dcim/sites/')

    def get_site(self, **kwargs):
        """Get site by filter

        :param kwargs: Filter can be any field.
        :return: site result
        """
        return self.netbox_con.get('/dcim/sites/', **kwargs)

    def create_site(self, name, slug, **kwargs):
        """Create new site

        :param name: Site name
        :param slug: slug name
        :param kwargs: optional fields
        :return: bool True if successful otherwise exception raised
        """
        required_fields = {"name": name, "slug": slug}
        return self.netbox_con.post('/dcim/sites/', required_fields, **kwargs)

    def delete_site(self, site_name):
        """Delete site

        :param site_name: Site to delete
        :return: bool True if succesful otherwase delete exception
        """
        site_id = self.get_site(name=site_name)[0]['id']
        return self.netbox_con.delete('/dcim/sites/', site_id)

    def get_racks(self):
        """Returns the available racks"""
        return self.netbox_con.get('/dcim/racks/')

    def get_rack(self, **kwargs):
        """Return rack id based on name

        :param kwargs:
        :return: Rack
        """
        return self.netbox_con.get('/dcim/racks/', **kwargs)

    def create_rack(self, name, site_name, **kwargs):
        """Create new rack

        :param name: Organizational rack name
        :param site_name: The site at which the rack exists
        :param kwargs: Optional arguments
        :return: bool True if successful otherwise create exception
        """
        site_id = self.get_site(name=site_name)[0]['id']
        required_fields = {"name": name, "site": site_id}
        return self.netbox_con.post('/dcim/racks/', required_fields, **kwargs)

    def delete_rack(self, rack_name):
        """Delete rack

        :param rack_name: Name of the rack to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        rack_id = self.get_rack(facility_id=rack_name)[0]['id']
        return self.netbox_con.delete('/dcim/racks/', rack_id)

    def get_devices(self):
        """Get all devices"""
        return self.netbox_con.get('/dcim/devices/')

    def get_device(self, **kwargs):
        """Get device by filter

        :param kwargs: Filter can be any field.
        :return: device result
        """
        return self.netbox_con.get('/dcim/devices/', **kwargs)

    def get_devices_per_rack(self, rack_name):
        """Get devices which belongs to the given rack

        :param rack_name: Name of the rack
        :return: list of devices otherwise an empty list
        """
        rack_id = self.get_rack(facility_id=rack_name)[0]['id']
        return self.netbox_con.get('/dcim/devices', rack_id=rack_id)

    def create_device(self, name, device_role, site_name, device_type, **kwargs):
        """Create a new device

        :param name: Name of the device
        :param device_role: Device role for the device
        :param site_name: Name of the site where the device is created
        :param device_type: Type for the new device
        :param kwargs: Optional arguments
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"name": name}
        device_role_id = self.get_device_role(name=device_role)[0]['id']
        site_id = self.get_site(name=site_name)[0]['id']
        device_type_id = self.get_device_type(model=device_type)[0]['id']

        if device_role_id:
            required_fields.update({"device_role": device_role_id})
        else:
            err_msg = 'Unable to create device. device role {} not found'.format(device_role)
            raise ValueError(err_msg)

        if site_id:
            required_fields.update({"site": site_id})
        else:
            err_msg = 'Unable to create device. site {} not found'.format(site_name)
            raise ValueError(err_msg)

        if device_type_id:
            required_fields.update({"device_type": device_type_id})
        else:
            err_msg = 'Unable to create device. device_type {} not found'.format(device_type)
            raise ValueError(err_msg)

        return self.netbox_con.post('/dcim/devices/', required_fields, **kwargs)

    def delete_device(self, device_name):
        """

        :param device_name: Device to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        device_id = self.get_device(name=device_name)[0]['id']
        return self.netbox_con.delete('/dcim/devices/', device_id)

    def update_device(self, device, **kwargs):
        """

        :param device: device name to update
        :param kwargs: requests body dict
        :return: bool True if successful otherwise raise UpdateException
        """
        device_id = self.get_device(name=device)[0]['id']
        return self.netbox_con.patch('/dcim/devices/', device_id, **kwargs)

    def get_device_types(self):
        """Get devices by device type"""
        return self.netbox_con.get('/dcim/device-types/')

    def get_device_type(self, **kwargs):
        """Get device type by filter

        :param kwargs: Search filter
        :return: device type
        """
        return self.netbox_con.get('/dcim/device-types/', **kwargs)

    def create_device_type(self, model, slug, manufacturer, **kwargs):
        """Create device type

        :param model: Model name
        :param slug: Slug name
        :param manufacturer: Name of the manufacurer
        :param kwargs: optional arguments
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"model": model, "slug": slug, "manufacturer": manufacturer}
        return self.netbox_con.post('/dcim/device-types/', required_fields, **kwargs)

    def delete_device_type(self, model_name):
        """Delete device type

        :param model_name: Name of the model
        :return: bool True if successful otherwise raise DeleteException
        """
        device_type_id = self.get_device_type(model=model_name)[0]['id']
        return self.netbox_con.delete('/dcim/device-types/', device_type_id)

    def get_device_roles(self):
        """Return all the device roles"""
        return self.netbox_con.get('/dcim/device-roles/')

    def get_device_role(self, **kwargs):
        """Get device role by filter

        :param kwargs: Search filter
        :return: device role
        """
        return self.netbox_con.get('/dcim/device-roles/', **kwargs)

    def create_device_role(self, name, color, slug, **kwargs):
        """Create device role

        :param name: Role name
        :param color: HTML color code
        :param slug: Slug name
        :param kwargs: optional arguments
        :return: bool True if successful otherwise CreateException
        """
        required_fields = {"name": name, "color": color, "slug": slug}
        return self.netbox_con.post('/dcim/device-roles/', required_fields, **kwargs)

    def delete_device_role(self, device_role):
        """Delete device by device role

        :param device_role: name of the role
        :return: bool True if successful otherwise raise DeleteException
        """
        device_role_id = self.get_device_role(name=device_role)[0]['id']
        return self.netbox_con.delete('/dcim/device-roles/', device_role_id)

    def get_manufactures(self):
        """Return all manufactures"""
        return self.netbox_con.get('/dcim/manufacturers/')

    def get_manufacturer(self, **kwargs):
        """Get manufacturer by filter

        :param kwargs: Search filter
        :return: manufacturer
        """
        return self.netbox_con.get('/dcim/manufacturers/', **kwargs)

    def create_manufacturer(self, name, slug, **kwargs):
        """Create new manufacturer

        :param name: Name of manufacturer
        :param slug: Name of slug
        :param kwargs: Optional arguments
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"name": name, "slug": slug}
        return self.netbox_con.post('/dcim/manufacturers/', required_fields, **kwargs)

    def delete_manufacturer(self, manufacturer_name):
        """Delete manufacturer

        :param manufacturer_name: Name of manufacturer to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        manufacturer_id = self.get_manufacturer(name=manufacturer_name)[0]['id']
        return self.netbox_con.delete('/dcim/manufacturers/', manufacturer_id)

    def get_platforms(self):
        """Return all platforms"""
        return self.netbox_con.get('/dcim/platforms')

    def get_platform(self, **kwargs):
        """Get platform by filter

        :param kwargs: Search filter
        :return: platform
        """
        return self.netbox_con.get('/dcim/platform/', **kwargs)

    def create_platform(self, name, slug, **kwargs):
        """Create new platform

        :param name: Name of platform
        :param slug: Name of slug
        :param kwargs: Optional arguments
        :return:
        """
        required_fields = {"name": name, "slug": slug}
        return self.netbox_con.post('/dcim/platforms/', required_fields, **kwargs)

    def delete_platform(self, platform_name):
        """Delete platform

        :param platform_name: Name of platform to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        platform_id = self.get_platform(name=platform_name)[0]['id']
        return self.netbox_con.delete('/dcim/platforms/', platform_id)

    def get_interfaces(self):
        """Return interfaces"""
        return self.netbox_con.get('/dcim/interfaces')

    def get_interface(self, **kwargs):
        """Get interface by filter

        :param kwargs: Search filter
        :return: interface
        """
        return self.netbox_con.get('/dcim/interfaces/', **kwargs)

    def get_interfaces_by_device(self, device_name):
        """Get all interfaces by device

        :param device_name: Name of device to get interfaces off
        :return: list of interfaces
        """
        return self.netbox_con.get('/dcim/interfaces', device=device_name)

    def create_interface(self, name, form_facter, device, **kwargs):
        """Create a new interface

        :param name: name of the interface
        :param form_facter: interface type. It is not possible to get the list of form factors from the api. Search
        in the netbox code for the correct form factor number.
        :param kwargs: optional arguments
        :param device: name of the device to associate interface with
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"name": name, "form_factor": form_facter, "device": device}
        return self.netbox_con.post('/dcim/interfaces/', required_fields, **kwargs)

    def delete_interface_by_id(self, interface_id):
        """Delete interface by id

        :param interface_id: id of interface to remove
        :return:
        """
        return self.netbox_con.delete('/dcim/interfaces/', interface_id)

<commit_msg>implemented interface-connections<commit_after>class Dcim(object):

    def __init__(self, netbox_con):
        self.netbox_con = netbox_con

    def get_sites(self):
        """Returns the available sites"""
        return self.netbox_con.get('/dcim/sites/')

    def get_site(self, **kwargs):
        """Get site by filter

        :param kwargs: Filter can be any field.
        :return: site result
        """
        return self.netbox_con.get('/dcim/sites/', **kwargs)

    def create_site(self, name, slug, **kwargs):
        """Create new site

        :param name: Site name
        :param slug: slug name
        :param kwargs: optional fields
        :return: bool True if successful otherwise exception raised
        """
        required_fields = {"name": name, "slug": slug}
        return self.netbox_con.post('/dcim/sites/', required_fields, **kwargs)

    def delete_site(self, site_name):
        """Delete site

        :param site_name: Site to delete
        :return: bool True if succesful otherwase delete exception
        """
        site_id = self.get_site(name=site_name)[0]['id']
        return self.netbox_con.delete('/dcim/sites/', site_id)

    def get_racks(self):
        """Returns the available racks"""
        return self.netbox_con.get('/dcim/racks/')

    def get_rack(self, **kwargs):
        """Return rack id based on name

        :param kwargs:
        :return: Rack
        """
        return self.netbox_con.get('/dcim/racks/', **kwargs)

    def create_rack(self, name, site_name, **kwargs):
        """Create new rack

        :param name: Organizational rack name
        :param site_name: The site at which the rack exists
        :param kwargs: Optional arguments
        :return: bool True if successful otherwise create exception
        """
        site_id = self.get_site(name=site_name)[0]['id']
        required_fields = {"name": name, "site": site_id}
        return self.netbox_con.post('/dcim/racks/', required_fields, **kwargs)

    def delete_rack(self, rack_name):
        """Delete rack

        :param rack_name: Name of the rack to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        rack_id = self.get_rack(facility_id=rack_name)[0]['id']
        return self.netbox_con.delete('/dcim/racks/', rack_id)

    def get_devices(self):
        """Get all devices"""
        return self.netbox_con.get('/dcim/devices/')

    def get_device(self, **kwargs):
        """Get device by filter

        :param kwargs: Filter can be any field.
        :return: device result
        """
        return self.netbox_con.get('/dcim/devices/', **kwargs)

    def get_devices_per_rack(self, rack_name):
        """Get devices which belongs to the given rack

        :param rack_name: Name of the rack
        :return: list of devices otherwise an empty list
        """
        rack_id = self.get_rack(facility_id=rack_name)[0]['id']
        return self.netbox_con.get('/dcim/devices', rack_id=rack_id)

    def create_device(self, name, device_role, site_name, device_type, **kwargs):
        """Create a new device

        :param name: Name of the device
        :param device_role: Device role for the device
        :param site_name: Name of the site where the device is created
        :param device_type: Type for the new device
        :param kwargs: Optional arguments
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"name": name}
        device_role_id = self.get_device_role(name=device_role)[0]['id']
        site_id = self.get_site(name=site_name)[0]['id']
        device_type_id = self.get_device_type(model=device_type)[0]['id']

        if device_role_id:
            required_fields.update({"device_role": device_role_id})
        else:
            err_msg = 'Unable to create device. device role {} not found'.format(device_role)
            raise ValueError(err_msg)

        if site_id:
            required_fields.update({"site": site_id})
        else:
            err_msg = 'Unable to create device. site {} not found'.format(site_name)
            raise ValueError(err_msg)

        if device_type_id:
            required_fields.update({"device_type": device_type_id})
        else:
            err_msg = 'Unable to create device. device_type {} not found'.format(device_type)
            raise ValueError(err_msg)

        return self.netbox_con.post('/dcim/devices/', required_fields, **kwargs)

    def delete_device(self, device_name):
        """

        :param device_name: Device to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        device_id = self.get_device(name=device_name)[0]['id']
        return self.netbox_con.delete('/dcim/devices/', device_id)

    def update_device(self, device, **kwargs):
        """

        :param device: device name to update
        :param kwargs: requests body dict
        :return: bool True if successful otherwise raise UpdateException
        """
        device_id = self.get_device(name=device)[0]['id']
        return self.netbox_con.patch('/dcim/devices/', device_id, **kwargs)

    def get_device_types(self):
        """Get devices by device type"""
        return self.netbox_con.get('/dcim/device-types/')

    def get_device_type(self, **kwargs):
        """Get device type by filter

        :param kwargs: Search filter
        :return: device type
        """
        return self.netbox_con.get('/dcim/device-types/', **kwargs)

    def create_device_type(self, model, slug, manufacturer, **kwargs):
        """Create device type

        :param model: Model name
        :param slug: Slug name
        :param manufacturer: Name of the manufacurer
        :param kwargs: optional arguments
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"model": model, "slug": slug, "manufacturer": manufacturer}
        return self.netbox_con.post('/dcim/device-types/', required_fields, **kwargs)

    def delete_device_type(self, model_name):
        """Delete device type

        :param model_name: Name of the model
        :return: bool True if successful otherwise raise DeleteException
        """
        device_type_id = self.get_device_type(model=model_name)[0]['id']
        return self.netbox_con.delete('/dcim/device-types/', device_type_id)

    def get_device_roles(self):
        """Return all the device roles"""
        return self.netbox_con.get('/dcim/device-roles/')

    def get_device_role(self, **kwargs):
        """Get device role by filter

        :param kwargs: Search filter
        :return: device role
        """
        return self.netbox_con.get('/dcim/device-roles/', **kwargs)

    def create_device_role(self, name, color, slug, **kwargs):
        """Create device role

        :param name: Role name
        :param color: HTML color code
        :param slug: Slug name
        :param kwargs: optional arguments
        :return: bool True if successful otherwise CreateException
        """
        required_fields = {"name": name, "color": color, "slug": slug}
        return self.netbox_con.post('/dcim/device-roles/', required_fields, **kwargs)

    def delete_device_role(self, device_role):
        """Delete device by device role

        :param device_role: name of the role
        :return: bool True if successful otherwise raise DeleteException
        """
        device_role_id = self.get_device_role(name=device_role)[0]['id']
        return self.netbox_con.delete('/dcim/device-roles/', device_role_id)

    def get_manufactures(self):
        """Return all manufactures"""
        return self.netbox_con.get('/dcim/manufacturers/')

    def get_manufacturer(self, **kwargs):
        """Get manufacturer by filter

        :param kwargs: Search filter
        :return: manufacturer
        """
        return self.netbox_con.get('/dcim/manufacturers/', **kwargs)

    def create_manufacturer(self, name, slug, **kwargs):
        """Create new manufacturer

        :param name: Name of manufacturer
        :param slug: Name of slug
        :param kwargs: Optional arguments
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"name": name, "slug": slug}
        return self.netbox_con.post('/dcim/manufacturers/', required_fields, **kwargs)

    def delete_manufacturer(self, manufacturer_name):
        """Delete manufacturer

        :param manufacturer_name: Name of manufacturer to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        manufacturer_id = self.get_manufacturer(name=manufacturer_name)[0]['id']
        return self.netbox_con.delete('/dcim/manufacturers/', manufacturer_id)

    def get_platforms(self):
        """Return all platforms"""
        return self.netbox_con.get('/dcim/platforms')

    def get_platform(self, **kwargs):
        """Get platform by filter

        :param kwargs: Search filter
        :return: platform
        """
        return self.netbox_con.get('/dcim/platform/', **kwargs)

    def create_platform(self, name, slug, **kwargs):
        """Create new platform

        :param name: Name of platform
        :param slug: Name of slug
        :param kwargs: Optional arguments
        :return:
        """
        required_fields = {"name": name, "slug": slug}
        return self.netbox_con.post('/dcim/platforms/', required_fields, **kwargs)

    def delete_platform(self, platform_name):
        """Delete platform

        :param platform_name: Name of platform to delete
        :return: bool True if successful otherwise raise DeleteException
        """
        platform_id = self.get_platform(name=platform_name)[0]['id']
        return self.netbox_con.delete('/dcim/platforms/', platform_id)

    def get_interfaces(self):
        """Return interfaces"""
        return self.netbox_con.get('/dcim/interfaces')

    def get_interface(self, **kwargs):
        """Get interface by filter

        :param kwargs: Search filter
        :return: interface
        """
        return self.netbox_con.get('/dcim/interfaces/', **kwargs)

    def get_interfaces_by_device(self, device_name):
        """Get all interfaces by device

        :param device_name: Name of device to get interfaces off
        :return: list of interfaces
        """
        return self.netbox_con.get('/dcim/interfaces', device=device_name)

    def create_interface(self, name, form_facter, device, **kwargs):
        """Create a new interface

        :param name: name of the interface
        :param form_facter: interface type. It is not possible to get the list of form factors from the api. Search
        in the netbox code for the correct form factor number.
        :param kwargs: optional arguments
        :param device: name of the device to associate interface with
        :return: bool True if successful otherwise raise CreateException
        """
        required_fields = {"name": name, "form_factor": form_facter, "device": device}
        return self.netbox_con.post('/dcim/interfaces/', required_fields, **kwargs)

    def delete_interface_by_id(self, interface_id):
        """Delete interface by id

        :param interface_id: id of interface to remove
        :return:
        """
        return self.netbox_con.delete('/dcim/interfaces/', interface_id)

    def get_interface_connections(self, **kwargs):
        """Get interface connections

        :param kwargs: Filter arguments
        :return: list of interface connections
        """
        return self.netbox_con.get('/dcim/interface-connections/', **kwargs)

    def create_interface_connection(self, interface_a, interface_b, **kwargs):
        """Create a new interface-connection

        :param interface_a: id of the source interface
        :param interface_b: id of the destination interface
        :param kwargs:
        :return:
        """
        required_fields = {"interface_a": interface_a, "interface_b": interface_b}
        return self.netbox_con.post('/dcim/interface-connections/', required_fields, **kwargs)

<|endoftext|>


<commit_before>""" A Data Access Object.

Contains the logic for inserting, updating, listing object in
a SQL database.
"""

import db
import json

from protorpc import message_types
from protorpc import messages
from protorpc import protojson

# Types that should be represented in a query format with '%s'
STRING_TYPES = (
    messages.StringField,
    message_types.DateTimeField,
)

# Types that should be represented in a query format with '%d'
NUMERIC_TYPES = (
    messages.BooleanField,
    messages.EnumField,
    messages.IntegerField,
)

class DbException(BaseException):
  """Exception thrown when a DB error is encounterd"""

class MissingKeyException(BaseException):
  """Exception thrown when a Get request doesn't conain all required keys."""

class NotFoundException(BaseException):
  """Exception thrown when a db object is not found."""

class DataAccessObject(object):
  """A DataAccessObject handles the mapping of object to the datbase.

  """
  def __init__(self, resource, table, columns, key_columns, column_map=None):
    """Constructs a DataAccessObject.

    Note: If a field in the resource contains an optionally repeating child
    object, it can be persisted two ways.  By specifying the field in the list
    of columns, the child object(s) will be converted to JSON and stored in that
    column.  If the child object is to be stored in a seperate table, the field
    should not be in the list of columns.  Instead, a DAO should be created for
    the child object type, and add_child_message function should be called to
    register it.  In addition, a link() should be overridden on the child DAO,
    and assemble() should be overridden on the parent DAO.

    Args:
    resource: The resource object. (The object containing the data).
    table: The table that these objects are stored in.
    columns: The columns of the database table that map to fields that should
        be persisted. The names of the columns should match the columns.  If
        that is not possible, pass in a column map.
    key_columns: The columns that are part of the object's primary key.
    column_map: A dictionary containing a map of column names to field values.
        Entries in this map are only necessary if they differ.  If all fields
        match the column names, this can be omitted.
    """
    self.resource = resource
    self.table = table
    self.columns = columns
    self.primary_key = _PrimaryKey(resource, key_columns)
    self.column_map = column_map or {}
    self.children = []

  def add_child_message(self, field_name, dao):
    """Registers a field as containing child messages.

    Args:
      field_name: The name of the field in the resource.
      dao: An instance of the Data Access Object that is responsible for the
          child object.
    """
    self.children.append((field_name, dao))

  def insert(self, obj):
    """Inserts this object into the database"""
    return self._insert_or_update(obj, update=False)

  def update(self, obj):
    return self._insert_or_update(obj, update=True)

  def get(self, request_obj):
    where_clause, keys = self.primary_key.where_clause(request_obj)
    if len(keys) != len(self.primary_key.columns()):
      raise MissingKeyException('Get {} requires {} to be specified'.format(
          self.table, self.primary_key.columns()))

    results = self._query(where_clause, keys)
    if not results:
      raise NotFoundException("Object not found. {} {}".format(where_clause,
                                                               keys))
    if len(results) == 0:
      raise DbException("Get returned {} results. {} {}".format(
          len(results), where_clause, keys))
    return results[0]

  def list(self, request_obj):
    return self._query(*self.primary_key.where_clause(request_obj))

  @db.connection
  def _insert_or_update(self, connection, obj, update=False):
    self._insert_update_with_connection(connection, obj, update)
    connection.commit()
    return self.get(obj)

  def _insert_update_with_connection(self, connection, obj, update):
    placeholders = []
    vals = []
    cols = []
    for col in self.columns:
      field_name = self._column_to_field(col)
      field = getattr(self.resource, field_name)
      val = getattr(obj, field_name)
      # Only use values that are present, and don't update the primary keys.
      if val is not None and not (update and col in self.primary_key.columns()):
        cols.append(col)
        placeholders.append(_placeholder_for_type(type(field), col))
        vals.append(_marshall_field(field, val))

    where_clause, keys = self.primary_key.where_clause(obj)
    if update:
      vals += keys
      assignments = ','.join(
          '{}={}'.format(k, v) for k, v in zip(cols, placeholders))
      query = 'UPDATE {} SET {} {}'.format(self.table, assignments,
                                           where_clause)
    else:
      query = 'INSERT INTO {} ({}) VALUES ({})'.format(
          self.table, ','.join(cols), ','.join(placeholders))
      connection.cursor().execute(query, vals)

    for field, dao in self.children:
      val = getattr(obj, field, None)
      if val is not None:
        field_def = getattr(self.resource, field)
        if field_def.repeated:
          children = val
        else:
          children = [val]
        for i, child in enumerate(children):
          dao.link(child, obj, ordinal=i)
          dao._insert_update_with_connection(connection, child, update)

  def link(self, obj, parent, ordinal):
    """Override this to propagate information across an object hierarchy.

    The default implementation does nothing.

    This is called once for every child object in an object hierarchy before
    insert or update.  This can be used to set the parent_id field of a child
    object to match parent.id.

    This can also be used to set ordinal fields to ensure that child object
    order is preserved.

    Args:
      obj: The child object.
      parent: The parent of this child object.
      ordinal: The order of this child object within the parent.
    """
    pass

  def assemble(self, obj):
    """Override this on the DAO to assemble the object hierarchy.

    Override this only for the root object in the hierarchy.

    This function should query all the child objects for the entire hierarchy
    and assemble them.  It is also the responsibility of this object to ensure
    that all child objects appear in the correct order.

    Args:
     obj: The root object.
    """
    pass

  @db.connection
  def _query(self, connection, where='', where_vals=None):
    """Loads object that match the given where clause and values.

    If no where clause is specified, all objects are loaded.
    """
    where_vals = where_vals or []
    query = 'SELECT {} from {} {}'.format(','.join(self.columns), self.table,
                                          where)
    cursor = connection.cursor()
    cursor.execute(query, where_vals)
    results = cursor.fetchall()
    objs = []
    for result in results:
      obj = self.resource()
      for i, col in enumerate(self.columns):
        field_name = self._column_to_field(col)
        field = getattr(self.resource, field_name)
        setattr(obj, field_name, _unmarshall_field(field, result[i]))

      self.assemble(obj)
      objs.append(obj)
    return objs

  def _column_to_field(self, col):
    """Maps the column name to the resource field."""
    return self.column_map.get(col, col)

class _PrimaryKey(object):
  """A potentially multi-column primary key."""
  def __init__(self, resource, keys):
    self.key_columns = keys
    self.resource = resource

  def where_clause(self, obj):
    """Builds a where clause for this resource.

    Args:
      obj: If specified, the generated where clause will only contain
          fields present in this object.

    Returns:
     A tuple containing a where clause with placeholders and the values
      for those placeholders.
    """
    placeholders = []
    cols = []
    keys = []
    for col in self.key_columns:
      val = getattr(obj, col, None)
      if val:
        field = getattr(self.resource, col)
        placeholders.append(_placeholder_for_type(type(field), col))
        cols.append(col)
        keys.append(_marshall_field(field, val))

    if len(cols):
      clause = 'where ' + ' and '.join(
          '{}={}'.format(k, v) for k, v in zip(cols, placeholders))
    else:
      clause = ''

    return (clause, keys)

  def columns(self):
    return self.key_columns

def _placeholder_for_type(field_type, field_name):
  """Returns: The proper placeholder for the given field."""

  if field_type in STRING_TYPES:
    return '%s'
  elif field_type in NUMERIC_TYPES:
    return '%d'
  elif field_type == messages.MessageField:
    return '%s'
  else:
    raise DbException(
        'Can\'t handle type: {} {}'.format(field_type, field_name))

def _marshall_field(field, val):
  """Converts a field value to db representation.

  For most objects it just returns the object. For Enums, for example, it
  converts to an int.
  """
  field_type = type(field)
  if field_type == messages.EnumField or field_type == messages.BooleanField:
    return int(val)

  # If we are saving a message to a column, convert to JSON.
  # Json encoder doesn't do protos. Proto encoder doesn't do arrays.
  if field_type == messages.MessageField:
    if field.repeated:
      return json.dumps([json.loads(protojson.encode_message(m)) for m in val])
    else:
      return protojson.encode_message(val)
  return val

def _unmarshall_field(field, result):
  if result == None:
    return None

  if type(field) == messages.EnumField:
    return field.type(result)
  elif type(field) == messages.MessageField:
    if field.repeated:
      # Json decoder doesn't do protos. Proto decoder doesn't do arrays. Use
      # json to parse the array, dumps() each element, then parse as proto.
      return [protojson.decode_message(field.message_type, json.dumps(sub_obj))
              for sub_obj in json.loads(result)]
    else:
      return protojson.decode_message(field.message_type, result)
  elif type(field) == messages.BooleanField:
    return bool(result)
  else:
    return result
<commit_msg>insert_update -> insert_or_update for clarity.<commit_after>""" A Data Access Object.

Contains the logic for inserting, updating, listing object in
a SQL database.
"""

import db
import json

from protorpc import message_types
from protorpc import messages
from protorpc import protojson

# Types that should be represented in a query format with '%s'
STRING_TYPES = (
    messages.StringField,
    message_types.DateTimeField,
)

# Types that should be represented in a query format with '%d'
NUMERIC_TYPES = (
    messages.BooleanField,
    messages.EnumField,
    messages.IntegerField,
)

class DbException(BaseException):
  """Exception thrown when a DB error is encounterd"""

class MissingKeyException(BaseException):
  """Exception thrown when a Get request doesn't conain all required keys."""

class NotFoundException(BaseException):
  """Exception thrown when a db object is not found."""

class DataAccessObject(object):
  """A DataAccessObject handles the mapping of object to the datbase.

  """
  def __init__(self, resource, table, columns, key_columns, column_map=None):
    """Constructs a DataAccessObject.

    Note: If a field in the resource contains an optionally repeating child
    object, it can be persisted two ways.  By specifying the field in the list
    of columns, the child object(s) will be converted to JSON and stored in that
    column.  If the child object is to be stored in a seperate table, the field
    should not be in the list of columns.  Instead, a DAO should be created for
    the child object type, and add_child_message function should be called to
    register it.  In addition, a link() should be overridden on the child DAO,
    and assemble() should be overridden on the parent DAO.

    Args:
    resource: The resource object. (The object containing the data).
    table: The table that these objects are stored in.
    columns: The columns of the database table that map to fields that should
        be persisted. The names of the columns should match the columns.  If
        that is not possible, pass in a column map.
    key_columns: The columns that are part of the object's primary key.
    column_map: A dictionary containing a map of column names to field values.
        Entries in this map are only necessary if they differ.  If all fields
        match the column names, this can be omitted.
    """
    self.resource = resource
    self.table = table
    self.columns = columns
    self.primary_key = _PrimaryKey(resource, key_columns)
    self.column_map = column_map or {}
    self.children = []

  def add_child_message(self, field_name, dao):
    """Registers a field as containing child messages.

    Args:
      field_name: The name of the field in the resource.
      dao: An instance of the Data Access Object that is responsible for the
          child object.
    """
    self.children.append((field_name, dao))

  def insert(self, obj):
    """Inserts this object into the database"""
    return self._insert_or_update(obj, update=False)

  def update(self, obj):
    return self._insert_or_update(obj, update=True)

  def get(self, request_obj):
    where_clause, keys = self.primary_key.where_clause(request_obj)
    if len(keys) != len(self.primary_key.columns()):
      raise MissingKeyException('Get {} requires {} to be specified'.format(
          self.table, self.primary_key.columns()))

    results = self._query(where_clause, keys)
    if not results:
      raise NotFoundException("Object not found. {} {}".format(where_clause,
                                                               keys))
    if len(results) == 0:
      raise DbException("Get returned {} results. {} {}".format(
          len(results), where_clause, keys))
    return results[0]

  def list(self, request_obj):
    return self._query(*self.primary_key.where_clause(request_obj))

  @db.connection
  def _insert_or_update(self, connection, obj, update=False):
    self._insert_or_update_with_connection(connection, obj, update)
    connection.commit()
    return self.get(obj)

  def _insert_or_update_with_connection(self, connection, obj, update):
    placeholders = []
    vals = []
    cols = []
    for col in self.columns:
      field_name = self._column_to_field(col)
      field = getattr(self.resource, field_name)
      val = getattr(obj, field_name)
      # Only use values that are present, and don't update the primary keys.
      if val is not None and not (update and col in self.primary_key.columns()):
        cols.append(col)
        placeholders.append(_placeholder_for_type(type(field), col))
        vals.append(_marshall_field(field, val))

    where_clause, keys = self.primary_key.where_clause(obj)
    if update:
      vals += keys
      assignments = ','.join(
          '{}={}'.format(k, v) for k, v in zip(cols, placeholders))
      query = 'UPDATE {} SET {} {}'.format(self.table, assignments,
                                           where_clause)
    else:
      query = 'INSERT INTO {} ({}) VALUES ({})'.format(
          self.table, ','.join(cols), ','.join(placeholders))
      connection.cursor().execute(query, vals)

    for field, dao in self.children:
      val = getattr(obj, field, None)
      if val is not None:
        field_def = getattr(self.resource, field)
        if field_def.repeated:
          children = val
        else:
          children = [val]
        for i, child in enumerate(children):
          dao.link(child, obj, ordinal=i)
          dao._insert_or_update_with_connection(connection, child, update)

  def link(self, obj, parent, ordinal):
    """Override this to propagate information across an object hierarchy.

    The default implementation does nothing.

    This is called once for every child object in an object hierarchy before
    insert or update.  This can be used to set the parent_id field of a child
    object to match parent.id.

    This can also be used to set ordinal fields to ensure that child object
    order is preserved.

    Args:
      obj: The child object.
      parent: The parent of this child object.
      ordinal: The order of this child object within the parent.
    """
    pass

  def assemble(self, obj):
    """Override this on the DAO to assemble the object hierarchy.

    Override this only for the root object in the hierarchy.

    This function should query all the child objects for the entire hierarchy
    and assemble them.  It is also the responsibility of this object to ensure
    that all child objects appear in the correct order.

    Args:
     obj: The root object.
    """
    pass

  @db.connection
  def _query(self, connection, where='', where_vals=None):
    """Loads object that match the given where clause and values.

    If no where clause is specified, all objects are loaded.
    """
    where_vals = where_vals or []
    query = 'SELECT {} from {} {}'.format(','.join(self.columns), self.table,
                                          where)
    cursor = connection.cursor()
    cursor.execute(query, where_vals)
    results = cursor.fetchall()
    objs = []
    for result in results:
      obj = self.resource()
      for i, col in enumerate(self.columns):
        field_name = self._column_to_field(col)
        field = getattr(self.resource, field_name)
        setattr(obj, field_name, _unmarshall_field(field, result[i]))

      self.assemble(obj)
      objs.append(obj)
    return objs

  def _column_to_field(self, col):
    """Maps the column name to the resource field."""
    return self.column_map.get(col, col)

class _PrimaryKey(object):
  """A potentially multi-column primary key."""
  def __init__(self, resource, keys):
    self.key_columns = keys
    self.resource = resource

  def where_clause(self, obj):
    """Builds a where clause for this resource.

    Args:
      obj: If specified, the generated where clause will only contain
          fields present in this object.

    Returns:
     A tuple containing a where clause with placeholders and the values
      for those placeholders.
    """
    placeholders = []
    cols = []
    keys = []
    for col in self.key_columns:
      val = getattr(obj, col, None)
      if val:
        field = getattr(self.resource, col)
        placeholders.append(_placeholder_for_type(type(field), col))
        cols.append(col)
        keys.append(_marshall_field(field, val))

    if len(cols):
      clause = 'where ' + ' and '.join(
          '{}={}'.format(k, v) for k, v in zip(cols, placeholders))
    else:
      clause = ''

    return (clause, keys)

  def columns(self):
    return self.key_columns

def _placeholder_for_type(field_type, field_name):
  """Returns: The proper placeholder for the given field."""

  if field_type in STRING_TYPES:
    return '%s'
  elif field_type in NUMERIC_TYPES:
    return '%d'
  elif field_type == messages.MessageField:
    return '%s'
  else:
    raise DbException(
        'Can\'t handle type: {} {}'.format(field_type, field_name))

def _marshall_field(field, val):
  """Converts a field value to db representation.

  For most objects it just returns the object. For Enums, for example, it
  converts to an int.
  """
  field_type = type(field)
  if field_type == messages.EnumField or field_type == messages.BooleanField:
    return int(val)

  # If we are saving a message to a column, convert to JSON.
  # Json encoder doesn't do protos. Proto encoder doesn't do arrays.
  if field_type == messages.MessageField:
    if field.repeated:
      return json.dumps([json.loads(protojson.encode_message(m)) for m in val])
    else:
      return protojson.encode_message(val)
  return val

def _unmarshall_field(field, result):
  if result == None:
    return None

  if type(field) == messages.EnumField:
    return field.type(result)
  elif type(field) == messages.MessageField:
    if field.repeated:
      # Json decoder doesn't do protos. Proto decoder doesn't do arrays. Use
      # json to parse the array, dumps() each element, then parse as proto.
      return [protojson.decode_message(field.message_type, json.dumps(sub_obj))
              for sub_obj in json.loads(result)]
    else:
      return protojson.decode_message(field.message_type, result)
  elif type(field) == messages.BooleanField:
    return bool(result)
  else:
    return result
<|endoftext|>


<commit_before>import logging
import requests

from indra.statements import Phosphorylation, Evidence, Agent

from .phosphoelm_mapping import phosphoelm_mapping

gilda_url = 'http://grounding.indra.bio/ground'
logger = logging.getLogger(__name__)

def _gilda_grounder(entity_str):
    # Try to find a namespace for the enzyme entity string, return the
    # string that provided the match
    res = requests.post(gilda_url, json={'text': entity_str})
    if res.status_code == 200 and res.json():
        db_ns = res.json()[0]['term']['db']
        db_id = res.json()[0]['term']['id']
        return entity_str, db_ns, db_id
    else:
        if res.status_code != 200:
            logger.warning('Gilda service responded with status code %d' %
                           res.status_code)
        return entity_str, None, None

class PhosphoElmProcessor(object):
    """Processes data dumps from the phospho.ELM database.

    See http://phospho.elm.eu.org/dataset.html

    Parameters
    ----------
    phosphoelm_data : list(dict)
        JSON compatible list of entries from a phospho.ELM data dump

    Attributes
    ----------
    statements : list(indra.statement.Phosphorylation)
        A list of the phosphorylation statements produced by the entries
        in phosphoelm_data
    """
    def __init__(self, phosphoelm_data):
        self.statements = []
        self._phosphoelm_data = phosphoelm_data

    def process_phosphorylations(self, keep_empty=False):
        """Create Phosphorylation statements from phosphoelm_data

        Parameters
        ----------
        keep_empty : bool
            Default: False. If true, also create statements when upstream
            kinases in entry['kinases'] are not known.
        """
        for entry in self._phosphoelm_data:
            if not keep_empty and not entry['kinases'] or\
                    not entry['species'].lower() == 'homo sapiens':
                # Skip entries without any kinases or if species is other
                # than human.
                continue
            # Entries:
            # 'acc': '<UP ID>', <-- substrate
            # 'sequence': '<protein sequence>',
            # 'position': '<sequence position>',
            # 'code': '<phosphorylated residue>',
            # 'pmids': '<pmid>',
            # 'kinases': '<responsible kinase>', <-- enzyme
            # 'source': 'HTP|LTP',
            # 'species': '<species name in latin>',
            # 'entry_date': 'yyyy-mm-dd HH:MM:SS.mmmmmm'
            substrate = Agent(None, db_refs={'UP': entry['acc']})
            used_name, enz = self._get_enzyme(entry['kinases']) if\
                entry.get('kinases') else '', None

            evidence = Evidence(
                source_api='phospho.ELM',
                pmid=entry['pmids'],
                annotations={
                    'data_source': 'High-ThroughPut' if
                    entry['source'].lower == 'htp' else (
                        'Low-ThroughPut' if entry['source'].lower() == 'ltp'
                        else None),
                    'phosphoelm_substrate_name': entry['acc'],
                    'phosphoelm_kinase_name': entry.get('kinases', None),
                    'used_kinse_name': used_name,
                    'entry_date': entry['entry_date'],
                    'sequence': entry['sequence']
                }
            )
            self.statements.append(Phosphorylation(
                enz=enz,
                sub=substrate,
                residue=entry['code'],
                position=entry['position'],
                evidence=evidence)
            )

    @staticmethod
    def _get_enzyme(upstream_kinase):
        """Handle the upstream kinases

        Parameters
        ----------
        upstream_kinase : str
            The string occuring in the entry 'upstream_kinases'

        Returns
        -------
        used_str, ag : tuple(str, indra.statements.Agent)
            A tuple containing a string and an agent. The string is the
            string from within 'upstream_kinases' that was used to create
            the agent.
        """
        # Check if kinase is in the hard coded mapping
        if upstream_kinase in phosphoelm_mapping:
            ns, _id = phosphoelm_mapping[upstream_kinase]
            return upstream_kinase, Agent(None, db_refs={ns: _id})

        strip_words = ['_group', '_drome', '_Caeel']
        # Pre process: strip 'strip words' and any trailing space
        for word in strip_words:
            upstream_kinase = upstream_kinase.replace(word, '').rstrip()
        used_str, ns, _id = _gilda_grounder(upstream_kinase)

        # Split on '/'
        if ns is None and _id is None and '/' in used_str:
            used_str = used_str.split('/')[0]
            used_str, ns, _id = _gilda_grounder(used_str)

        # Replace '_' with '-'
        if ns is None and _id is None and '_' in used_str:
            used_str = used_str.replace('_', '-')
            used_str, ns, _id = _gilda_grounder(used_str)

        if ns is None and _id is None:
            ns = 'TEXT'
            _id = used_str

        ag = Agent(None, db_refs={ns: _id})
        return used_str, ag
<commit_msg>Handle special cases when no Agent is returned<commit_after>import logging
import requests

from indra.statements import Phosphorylation, Evidence, Agent

from .phosphoelm_mapping import phosphoelm_mapping

gilda_url = 'http://grounding.indra.bio/ground'
logger = logging.getLogger(__name__)

def _gilda_grounder(entity_str):
    # Try to find a namespace for the enzyme entity string, return the
    # string that provided the match
    res = requests.post(gilda_url, json={'text': entity_str})
    if res.status_code == 200 and res.json():
        db_ns = res.json()[0]['term']['db']
        db_id = res.json()[0]['term']['id']
        return entity_str, db_ns, db_id
    else:
        if res.status_code != 200:
            logger.warning('Gilda service responded with status code %d' %
                           res.status_code)
        return entity_str, None, None

class PhosphoElmProcessor(object):
    """Processes data dumps from the phospho.ELM database.

    See http://phospho.elm.eu.org/dataset.html

    Parameters
    ----------
    phosphoelm_data : list(dict)
        JSON compatible list of entries from a phospho.ELM data dump

    Attributes
    ----------
    statements : list(indra.statement.Phosphorylation)
        A list of the phosphorylation statements produced by the entries
        in phosphoelm_data
    """
    def __init__(self, phosphoelm_data):
        self.statements = []
        self._phosphoelm_data = phosphoelm_data

    def process_phosphorylations(self, keep_empty=False):
        """Create Phosphorylation statements from phosphoelm_data

        Parameters
        ----------
        keep_empty : bool
            Default: False. If true, also create statements when upstream
            kinases in entry['kinases'] are not known.
        """
        for entry in self._phosphoelm_data:
            if not keep_empty and not entry['kinases'] or\
                    not entry['species'].lower() == 'homo sapiens':
                # Skip entries without any kinases or if species is other
                # than human.
                continue
            # Entries:
            # 'acc': '<UP ID>', <-- substrate
            # 'sequence': '<protein sequence>',
            # 'position': '<sequence position>',
            # 'code': '<phosphorylated residue>',
            # 'pmids': '<pmid>',
            # 'kinases': '<responsible kinase>', <-- enzyme
            # 'source': 'HTP|LTP',
            # 'species': '<species name in latin>',
            # 'entry_date': 'yyyy-mm-dd HH:MM:SS.mmmmmm'
            substrate = Agent(None, db_refs={'UP': entry['acc']})
            used_name, enz = self._get_enzyme(entry['kinases']) if\
                entry.get('kinases') else '', None

            # Skip if we hit one of the special cases
            if enz is None:
                continue

            # Build evidence, add statement
            evidence = Evidence(
                source_api='phospho.ELM',
                pmid=entry['pmids'],
                annotations={
                    'data_source': 'High-ThroughPut' if
                    entry['source'].lower == 'htp' else (
                        'Low-ThroughPut' if entry['source'].lower() == 'ltp'
                        else None),
                    'phosphoelm_substrate_name': entry['acc'],
                    'phosphoelm_kinase_name': entry.get('kinases', None),
                    'used_kinse_name': used_name,
                    'entry_date': entry['entry_date'],
                    'sequence': entry['sequence']
                }
            )
            self.statements.append(Phosphorylation(
                enz=enz,
                sub=substrate,
                residue=entry['code'],
                position=entry['position'],
                evidence=evidence)
            )

    @staticmethod
    def _get_enzyme(upstream_kinase):
        """Handle the upstream kinases

        Parameters
        ----------
        upstream_kinase : str
            The string occuring in the entry 'upstream_kinases'

        Returns
        -------
        used_str, ag : tuple(str, indra.statements.Agent)
            A tuple containing a string and an agent. The string is the
            string from within 'upstream_kinases' that was used to create
            the agent.
        """
        # Check if kinase is in the hard coded mapping
        if upstream_kinase in phosphoelm_mapping:
            ns, _id = phosphoelm_mapping[upstream_kinase]
            if ns is None and _id is None:
                return upstream_kinase, None
            return upstream_kinase, Agent(None, db_refs={ns: _id})

        strip_words = ['_group', '_drome', '_Caeel']
        # Pre process: strip 'strip words' and any trailing space
        for word in strip_words:
            upstream_kinase = upstream_kinase.replace(word, '').rstrip()
        used_str, ns, _id = _gilda_grounder(upstream_kinase)

        # Split on '/'
        if ns is None and _id is None and '/' in used_str:
            used_str = used_str.split('/')[0]
            used_str, ns, _id = _gilda_grounder(used_str)

        # Replace '_' with '-'
        if ns is None and _id is None and '_' in used_str:
            used_str = used_str.replace('_', '-')
            used_str, ns, _id = _gilda_grounder(used_str)

        if ns is None and _id is None:
            ns = 'TEXT'
            _id = used_str

        ag = Agent(None, db_refs={ns: _id})
        return used_str, ag
<|endoftext|>


<commit_before>#!/usr/bin/env python
from __future__ import print_function
from __future__ import unicode_literals

class Graph(object):
    ''' Create an empty graph. '''
    def __init__(self):
        self.graph = {}
        return

    def nodes():
        return nodes

    def edges():
        return edges

    def add_node(self, node):
        self.graph.setdefault(node, [])
        return

    def add_edge(self, node1, node2):
        return

    def del_node(self, node):
        try:
            del self.graph[node]
        except KeyError:
            raise KeyError('node not in graph')

    def has_node(self, node):
        return node in self.graph

    def neighbors(self, node):
        return self.graph[node]

    def adjecent(self, node1, node2):
        if node2 in self.graph[node1] or node1 in self.graph[node2]:
            return True
        else:
            return False
<commit_msg>Add functions for adding a node, an edge, and defining a node<commit_after>#!/usr/bin/env python
from __future__ import print_function
from __future__ import unicode_literals

class Graph(object):
    ''' Create an empty graph. '''
    def __init__(self):
        self.graph = {}
        return

    def nodes(self):
        return self.graph.keys()

    def edges(self):
        edge_list = []
        for key, value in self.graph():
            for item in value:
                edge_list.append((key, item))
        return edge_list

    def add_node(self, node):
        self.graph.setdefault(node, [])

    def add_edge(self, node1, node2):
        if node1 in self.graph:
            self.graph.append(node2)
        else:
            self.graph[node1] = node2

    def del_node(self, node):
        try:
            del self.graph[node]
        except KeyError:
            raise KeyError('node not in graph')

    def has_node(self, node):
        return node in self.graph

    def neighbors(self, node):
        return self.graph[node]

    def adjecent(self, node1, node2):
        if node2 in self.graph[node1] or node1 in self.graph[node2]:
            return True
        else:
            return False
<|endoftext|>


<commit_before><commit_msg>style: improve readability by applying black formatter<commit_after><|endoftext|>


<commit_before>import hashlib
import math
from os import path
from base64 import b64encode, b64decode
from Crypto.Cipher import DES, AES
from Crypto import Random

random_instance = Random.new()

algorithms = {'aes': AES, 'des': DES}

def sha1(message):
    return hashlib.sha1(message).hexdigest()

def md5(message):
    return hashlib.md5(message).hexdigest()

def sha256(message):
    return hashlib.sha256(message).hexdigest()

def sha512(message):
    return hashlib.sha152(message).hexdigest()

hashes = [sha1, md5, sha256, sha512]

def str_to_base64(message):
    return b64encode(message)

def base64_to_str(message):
    return b64decode(message)

def str_to_hex(message):
    return message.encode('hex')

def hex_to_str(message):
    return message.decode('hex')

base64 = str_to_base64
hex = str_to_hex
encodes = [str_to_base64, str_to_hex]
modifiers = [str.lower, str.upper, str.strip, str.rstrip,
             path.basename, path.abspath, path.dirname, path.relpath,
             lambda s: s + '\n', lambda s: '/' + s,
             lambda s: s.replace('\\', '/'), lambda s: s.replace('-', '')]
decodes = [base64_to_str, hex_to_str]

def apply_modifiers(modifiers, message):
    try:
        for modifier in modifiers:
            message = modifier(message)
    finally:
        return message

def guess_hash(message, hash_value, cur_guess=[]):
    if len(cur_guess) > 4:
        return None

    if apply_modifiers(cur_guess, message) == hash_value:
        return cur_guess

    for modifier in hashes + modifiers + decodes:
        next_guess = cur_guess + [modifier]
        if guess_hash(message, hash_value, next_guess):
            return next_guess

    return None

print guess_hash('teste\n', sha1('teste'))

def pad(message, length, padding=' '):
    return message + (length - len(message)) * padding

def pad_multiple(message, len_multiple, padding=' '):
    next_length = math.ceil(len(message) / float(len_multiple)) * len_multiple
    return pad(message, int(next_length), padding)

def random(n_bytes):
    return random_instance.read(n_bytes)

def encrypt(message, password, algorithm='aes'):
    cls = algorithms[algorithm]
    iv = random(cls.block_size)
    instance = cls.new(pad_multiple(password, 16),
                       cls.MODE_CFB,
                       iv)
    return str_to_base64(iv + instance.encrypt(message))

def decrypt(message, password, algorithm='aes'):
    message = base64_to_str(message)
    iv, message = message[:AES.block_size], message[AES.block_size:]
    instance = AES.new(pad_multiple(password, 16),
                       AES.MODE_CFB,
                       iv)
    return instance.decrypt(message)

def encrypt_aes(message, password):
    return encrypt(message, password, 'aes')

def decrypt_aes(message, password):
    return decrypt(message, password, 'aes')

def encrypt_des(message, password):
    return encrypt(message, password, 'des')

def decrypt_des(message, password):
    return decrypt(message, password, 'des')
<commit_msg>Use iteration instead of recursion in guessing<commit_after>import hashlib
import math
from os import path
from base64 import b64encode, b64decode
from Crypto.Cipher import DES, AES
from Crypto import Random

random_instance = Random.new()

algorithms = {'aes': AES, 'des': DES}

def sha1(message):
    return hashlib.sha1(message).hexdigest()

def md5(message):
    return hashlib.md5(message).hexdigest()

def sha256(message):
    return hashlib.sha256(message).hexdigest()

def sha512(message):
    return hashlib.sha152(message).hexdigest()

hashes = [sha1, md5, sha256, sha512]

def str_to_base64(message):
    return b64encode(message)

def base64_to_str(message):
    return b64decode(message)

def str_to_hex(message):
    return message.encode('hex')

def hex_to_str(message):
    return message.decode('hex')

def append_newline(s):
    return s + '\n'

def replace_backslashes(s):
    return s.replace('\\', '/')

base64 = str_to_base64
hex = str_to_hex
encodes = [str_to_base64, str_to_hex]
modifiers = [str.lower, str.upper, str.strip,
             append_newline, replace_backslashes,
             path.basename, path.abspath, path.dirname]
decodes = [base64_to_str, hex_to_str]

def apply_modifiers(modifiers, message):
    try:
        for modifier in modifiers:
            message = modifier(message)
    finally:
        return message

def guess_hash(message, hash_value):
    from collections import deque
    guesses = deque()
    guesses.append([])
    while len(guesses):
        guess = guesses.popleft()
        if apply_modifiers(guess, message) == hash_value:
            return guess

        if len(guess) < 5:
            for modifier in hashes + modifiers + decodes + encodes:
                guesses.append(guess + [modifier])
    return None

def pad(message, length, padding=' '):
    return message + (length - len(message)) * padding

def pad_multiple(message, len_multiple, padding=' '):
    next_length = math.ceil(len(message) / float(len_multiple)) * len_multiple
    return pad(message, int(next_length), padding)

def random(n_bytes):
    return random_instance.read(n_bytes)

def encrypt(message, password, algorithm='aes'):
    cls = algorithms[algorithm]
    iv = random(cls.block_size)
    instance = cls.new(pad_multiple(password, 16),
                       cls.MODE_CFB,
                       iv)
    return str_to_base64(iv + instance.encrypt(message))

def decrypt(message, password, algorithm='aes'):
    message = base64_to_str(message)
    iv, message = message[:AES.block_size], message[AES.block_size:]
    instance = AES.new(pad_multiple(password, 16),
                       AES.MODE_CFB,
                       iv)
    return instance.decrypt(message)

def encrypt_aes(message, password):
    return encrypt(message, password, 'aes')

def decrypt_aes(message, password):
    return decrypt(message, password, 'aes')

def encrypt_des(message, password):
    return encrypt(message, password, 'des')

def decrypt_des(message, password):
    return decrypt(message, password, 'des')
<|endoftext|>


<commit_before><commit_msg>Add back in ILOperations replace 'vars_referenced' with 'vars' added 'vars_used_in_address' fixed some class hierarchy issues<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ======================================================================
#
# hc.py
# ------
# Simple, pure Python model of the HC stream cipher with
# support for 128 and 256 bit keys.
#
# Note: Currently the model is only implementing HC-256.
#
#
# Author: Joachim Strömbergson
# Copyright (c) 2017, Assured AB
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or
# without modification, are permitted provided that the following
# conditions are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# ======================================================================

# ------------------------------------------------------------------
# Python module imports.
# ------------------------------------------------------------------
import sys

# -------------------------------------------------------------------
# HC
# The HC stream cipher class.
# ------------------------------------------------------------------
class HC():
    def __init__(self, verbose=False):
        # Constants
        self.MAX_W32 = 0xffffffff
        self.HC256_TSIZE = 1024
        self.HC128_TSIZE = 512
        self.DUMP_W_ELEMENTS = 64
        self.DUMP_PQ_ELEMENTS = 64

        # Allocate the internal state variables.
        self.verbose = verbose
        self.W = [0] * 2560
        self.P = [0] * 1024
        self.Q = [0] * 1024
        self.i = 0

    # Initialize the internal state given the key and iv.
    def init(self, key, iv):
        for i in range(8):
            self.W[i] = key[i]
            self.W[(i + 8)] = iv[i]

        for i in range(16, 2560):
            self.W[i] = (self.f2(self.W[(i - 2)]) + self.W[(i - 7)] +
                         self.f1(self.W[(i - 15)]) + self.W[(i - 16)] + i) & self.MAX_W32

        if self.verbose:
            self.dump_w()

        for i in range(1024):
            self.P[i] = self.W[(i + 512)]
            self.Q[i] = self.W[(i + 1536)]

        if self.verbose:
            print("P and Q before 4096 internal updates.")
            self.dump_pq()

        for i in range(4096):
            self.next()

        if self.verbose:
            print("P and Q after 4096 internal updates.")
            self.dump_pq()

    # Update internal state and return the next word.
    def next(self):
        j = self.i % 1024
        if (self.i % 2048) < 1024:
            if self.verbose:
                print("i = %04d, i3 = %04d, i10 = %04d, i1023 = %04d" %
                          (j, self.subm(j, 3), self.subm(j, 10), self.subm(j, 1023)))

                print("x[i] = 0x%08x, x[i3] = 0x%08x, x[i10] = 0x%08x, x[i1023] = 0x%08x" %
                          (self.P[j], self.P[self.subm(j, 3)],
                           self.P[self.subm(j, 10)], self.P[self.subm(j, 1023)]))

            self.P[j] = (self.P[j] + self.P[self.subm(j, 10)] +
                         self.g1(self.P[self.subm(j, 3)],
                         self.P[self.subm(j, 1023)])) & self.MAX_W32

            if self.verbose:
                print("New x[i] = 0x%08x" % self.P[j])
                print("")

            s = self.h1(self.P[self.subm(j, 12)]) ^ self.P[j]

        else:
            self.Q[j] = (self.Q[j] + self.Q[self.subm(j, 10)] +
                         self.g2(self.Q[self.subm(j, 3)],
                         self.Q[self.subm(j, 1023)])) & self.MAX_W32
            s = self.h2(self.Q[self.subm(j, 12)]) ^ self.Q[j]

        self.i += 1
        return s

    # Internal HC functions.
    def f1(self, x):
        return self.rotr(x, 7) ^ self.rotr(x, 18) ^ self.shr(x, 3)

    def f2(self, x):
        return self.rotr(x, 17) ^ self.rotr(x, 19) ^ self.shr(x, 10)

    def g1(self, x, y):
        qval = self.Q[((x ^ y) % 1024)]
        rot10 = self.rotr(x, 10)
        rot23 = self.rotr(y, 23)
        result = ((rot10 ^ rot23) + qval) & self.MAX_W32
        if self.verbose:
            print("In g1. x = 0x%08x, y = 0x%08x, res = 0x%08x" % (x, y, result))
            print("x ^ y = 0x%08x, q[x^y] = 0x%08x, rot10 = 0x%08x, rot23 = 0x%08x" %
                  ((x ^ y), qval, rot10, rot23))

        return result

    def g2(self, x, y):
        result = (self.rotr(x, 10) ^ self.rotr(y, 23) +
                  self.P[((x ^ y) % 1024)]) & self.MAX_W32
        if self.verbose:
            print("In g2. x = 0x%08x, y = 0x%08x, res = 0x%08x" % (x, y, result))
        return result

    def h1(self, x):
        (x0, x1, x2, x3) = self.w2b(x)
        return (self.Q[x0] + self.Q[(x1 + 256)] +
                self.Q[(x2 + 512)] + self.Q[(x3 + 768)]) & self.MAX_W32

    def h2(self, x):
        (x0, x1, x2, x3) = self.w2b(x)
        return (self.P[x0] + self.P[(x1 + 256)] +
                self.P[(x2 + 512)] + self.P[(x3 + 768)]) & self.MAX_W32

    # Helper functions needed to implement the HC functions.
    def rotr(self, w, b):
        return ((w >> b) | (w << (32 - b))) & self.MAX_W32

    def shr(self, w, b):
        return (w >> b)

    def w2b(self, x):
        x0 = x >> 24
        x1 = x >> 16 & 0xff
        x2 = x >> 8 & 0xff
        x3 = x & 0xff
        return (x0, x1, x2, x3)

    def subm(self, x, y):
        return (x - y) % 1024

    def dump_w(self):
        print("State of the W array:")
        for i in range(0, self.DUMP_W_ELEMENTS, 4):
            print("W[%04d..%04d]: 0x%08x 0x%08x 0x%08x 0x%08x" %
                      (i, i+3, self.W[i], self.W[i+1], self.W[i+2], self.W[i+3]))
        print("")

    def dump_pq(self):
        print("State of the P array:")
        for i in range(0, self.DUMP_PQ_ELEMENTS, 4):
            print("P[%04d..%04d]: 0x%08x 0x%08x 0x%08x 0x%08x" %
                      (i, i+3, self.P[i], self.P[i+1], self.P[i+2], self.P[i+3]))
        print("")
        print("State of the Q array:")
        for i in range(0, self.DUMP_PQ_ELEMENTS, 4):
            print("Q[%04d..%04d]: 0x%08x 0x%08x 0x%08x 0x%08x" %
                      (i, i+3, self.Q[i], self.Q[i+1], self.Q[i+2], self.Q[i+3]))
        print("")

# ------------------------------------------------------------------
# test_rotr()
# Test the hc help function rotr()
# ------------------------------------------------------------------
def test_rotr():
    my_hc = HC()
    print("Testing rotr help function.")
    for i in range(33):
        print("rotr %02d steps: 0x%08x" % (i, my_hc.rotr(0x00000001, i)))
    print("")

# ------------------------------------------------------------------
# test_hc()
# Test the HC implementation with 128 and 256 bit keys.
# ------------------------------------------------------------------
def test_hc():
    my_key = [0] * 8
    my_iv = [0] * 8

    my_hc = HC(verbose=True)
    my_hc.init(my_key, my_iv)

    for i in range(16):
        print("keystream %02d = 0x%08x" % (i, my_hc.next()))

# ------------------------------------------------------------------
# main()
# If executed tests the ChaCha class using known test vectors.
# ------------------------------------------------------------------
def main():
    print("Testing the HC stream cipher model")
    print("==================================")
    print

#    test_rotr()
    test_hc()

# ------------------------------------------------------------------
# __name__
# Python thingy which allows the file to be run standalone as
# well as parsed from within a Python interpreter.
# ------------------------------------------------------------------
if __name__ == "__main__":
    # Run the main function.
    sys.exit(main())

# ======================================================================
# EOF hc.py
# ======================================================================
<commit_msg>First iteration of P correct. Adding debug outputs for Q updates.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ======================================================================
#
# hc.py
# ------
# Simple, pure Python model of the HC stream cipher with
# support for 128 and 256 bit keys.
#
# Note: Currently the model is only implementing HC-256.
#
#
# Author: Joachim Strömbergson
# Copyright (c) 2017, Assured AB
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or
# without modification, are permitted provided that the following
# conditions are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# ======================================================================

# ------------------------------------------------------------------
# Python module imports.
# ------------------------------------------------------------------
import sys

# -------------------------------------------------------------------
# HC
# The HC stream cipher class.
# ------------------------------------------------------------------
class HC():
    def __init__(self, verbose=False):
        # Constants
        self.MAX_W32 = 0xffffffff
        self.HC256_TSIZE = 1024
        self.HC128_TSIZE = 512
        self.DUMP_W_ELEMENTS = 64
        self.DUMP_PQ_ELEMENTS = 64

        # Allocate the internal state variables.
        self.verbose = verbose
        self.W = [0] * 2560
        self.P = [0] * 1024
        self.Q = [0] * 1024
        self.i = 0

    # Initialize the internal state given the key and iv.
    def init(self, key, iv):
        for i in range(8):
            self.W[i] = key[i]
            self.W[(i + 8)] = iv[i]

        for i in range(16, 2560):
            self.W[i] = (self.f2(self.W[(i - 2)]) + self.W[(i - 7)] +
                         self.f1(self.W[(i - 15)]) + self.W[(i - 16)] + i) & self.MAX_W32

        if self.verbose:
            self.dump_w()

        for i in range(1024):
            self.P[i] = self.W[(i + 512)]
            self.Q[i] = self.W[(i + 1536)]

        if self.verbose:
            print("P and Q before 4096 internal updates.")
            self.dump_pq()

        for i in range(4096):
            self.next()

        if self.verbose:
            print("P and Q after 4096 internal updates.")
            self.dump_pq()

    # Update internal state and return the next word.
    def next(self):
        j = self.i % 1024
        if (self.i % 2048) < 1024:
            if self.verbose:
                print("i = %04d, i3 = %04d, i10 = %04d, i1023 = %04d" %
                          (j, self.subm(j, 3), self.subm(j, 10), self.subm(j, 1023)))

                print("In P: x[i] = 0x%08x, x[i3] = 0x%08x, x[i10] = 0x%08x, x[i1023] = 0x%08x" %
                          (self.P[j], self.P[self.subm(j, 3)],
                           self.P[self.subm(j, 10)], self.P[self.subm(j, 1023)]))

            self.P[j] = (self.P[j] + self.P[self.subm(j, 10)] +
                         self.g1(self.P[self.subm(j, 3)],
                         self.P[self.subm(j, 1023)])) & self.MAX_W32

            if self.verbose:
                print("New x[i] = 0x%08x" % self.P[j])
                print("")

            s = self.h1(self.P[self.subm(j, 12)]) ^ self.P[j]

        else:
            if self.verbose:
                print("i = %04d, i3 = %04d, i10 = %04d, i1023 = %04d" %
                          (j, self.subm(j, 3), self.subm(j, 10), self.subm(j, 1023)))

                print("In Q: x[i] = 0x%08x, x[i3] = 0x%08x, x[i10] = 0x%08x, x[i1023] = 0x%08x" %
                          (self.Q[j], self.Q[self.subm(j, 3)],
                           self.Q[self.subm(j, 10)], self.Q[self.subm(j, 1023)]))

            self.Q[j] = (self.Q[j] + self.Q[self.subm(j, 10)] +
                         self.g2(self.Q[self.subm(j, 3)],
                         self.Q[self.subm(j, 1023)])) & self.MAX_W32

            if self.verbose:
                print("New x[i] = 0x%08x" % self.Q[j])
                print("")

            s = self.h2(self.Q[self.subm(j, 12)]) ^ self.Q[j]

        self.i += 1
        return s

    # Internal HC functions.
    def f1(self, x):
        return self.rotr(x, 7) ^ self.rotr(x, 18) ^ self.shr(x, 3)

    def f2(self, x):
        return self.rotr(x, 17) ^ self.rotr(x, 19) ^ self.shr(x, 10)

    def g1(self, x, y):
        qval = self.Q[((x ^ y) % 1024)]
        rot10 = self.rotr(x, 10)
        rot23 = self.rotr(y, 23)
        result = ((rot10 ^ rot23) + qval) & self.MAX_W32
        if self.verbose:
            print("In g1. x = 0x%08x, y = 0x%08x, res = 0x%08x" % (x, y, result))
            print("x ^ y = 0x%08x, q[x^y] = 0x%08x, rot10 = 0x%08x, rot23 = 0x%08x" %
                  ((x ^ y), qval, rot10, rot23))

        return result

    def g2(self, x, y):
        result = (self.rotr(x, 10) ^ self.rotr(y, 23) +
                  self.P[((x ^ y) % 1024)]) & self.MAX_W32
        if self.verbose:
            print("In g2. x = 0x%08x, y = 0x%08x, res = 0x%08x" % (x, y, result))
        return result

    def h1(self, x):
        (x0, x1, x2, x3) = self.w2b(x)
        return (self.Q[x0] + self.Q[(x1 + 256)] +
                self.Q[(x2 + 512)] + self.Q[(x3 + 768)]) & self.MAX_W32

    def h2(self, x):
        (x0, x1, x2, x3) = self.w2b(x)
        return (self.P[x0] + self.P[(x1 + 256)] +
                self.P[(x2 + 512)] + self.P[(x3 + 768)]) & self.MAX_W32

    # Helper functions needed to implement the HC functions.
    def rotr(self, w, b):
        return ((w >> b) | (w << (32 - b))) & self.MAX_W32

    def shr(self, w, b):
        return (w >> b)

    def w2b(self, x):
        x0 = x >> 24
        x1 = x >> 16 & 0xff
        x2 = x >> 8 & 0xff
        x3 = x & 0xff
        return (x0, x1, x2, x3)

    def subm(self, x, y):
        return (x - y) % 1024

    def dump_w(self):
        print("State of the W array:")
        for i in range(0, self.DUMP_W_ELEMENTS, 4):
            print("W[%04d..%04d]: 0x%08x 0x%08x 0x%08x 0x%08x" %
                      (i, i+3, self.W[i], self.W[i+1], self.W[i+2], self.W[i+3]))
        print("")

    def dump_pq(self):
        print("State of the P array:")
        for i in range(0, self.DUMP_PQ_ELEMENTS, 4):
            print("P[%04d..%04d]: 0x%08x 0x%08x 0x%08x 0x%08x" %
                      (i, i+3, self.P[i], self.P[i+1], self.P[i+2], self.P[i+3]))
        print("")
        print("State of the Q array:")
        for i in range(0, self.DUMP_PQ_ELEMENTS, 4):
            print("Q[%04d..%04d]: 0x%08x 0x%08x 0x%08x 0x%08x" %
                      (i, i+3, self.Q[i], self.Q[i+1], self.Q[i+2], self.Q[i+3]))
        print("")

# ------------------------------------------------------------------
# test_rotr()
# Test the hc help function rotr()
# ------------------------------------------------------------------
def test_rotr():
    my_hc = HC()
    print("Testing rotr help function.")
    for i in range(33):
        print("rotr %02d steps: 0x%08x" % (i, my_hc.rotr(0x00000001, i)))
    print("")

# ------------------------------------------------------------------
# test_hc()
# Test the HC implementation with 128 and 256 bit keys.
# ------------------------------------------------------------------
def test_hc():
    my_key = [0] * 8
    my_iv = [0] * 8

    my_hc = HC(verbose=True)
    my_hc.init(my_key, my_iv)

    for i in range(16):
        print("keystream %02d = 0x%08x" % (i, my_hc.next()))

# ------------------------------------------------------------------
# main()
# If executed tests the ChaCha class using known test vectors.
# ------------------------------------------------------------------
def main():
    print("Testing the HC stream cipher model")
    print("==================================")
    print

#    test_rotr()
    test_hc()

# ------------------------------------------------------------------
# __name__
# Python thingy which allows the file to be run standalone as
# well as parsed from within a Python interpreter.
# ------------------------------------------------------------------
if __name__ == "__main__":
    # Run the main function.
    sys.exit(main())

# ======================================================================
# EOF hc.py
# ======================================================================
<|endoftext|>


<commit_before><commit_msg>fix a bug in parser detected by pylint -E<commit_after><|endoftext|>


<commit_before>__author__ = 'Milinda Perera'

import random as pyrandom
from time import time

from psi_3_round_paillier import PSI3RoundPaillier
from psi_2_round_paillier import PSI2RoundPaillier
from psi_2_round_elgamal import PSI2RoundElGamal

def timer(func, *pargs, **kargs):
    start = time()
    ret = func(*pargs, **kargs)
    elapsed = time() - start
    return elapsed, ret

def run_psi_3_round_paillier(server_set, client_set):
    psi = PSI3RoundPaillier(1024)
    server_out_1, server_state = psi.server_to_client_1(server_set)
    client_out_1 = psi.client_to_server(client_set, **server_out_1)
    server_out_2 = psi.server_to_client_2(client_out_1, **server_state)
    client_out_2 = psi.client_output(server_out_2)
    return client_out_2

def run_psi_2_round_paillier(server_set, client_set):
    psi = PSI2RoundPaillier(1024)
    client_out_1, client_state = psi.client_to_server(client_set)
    server_out = psi.server_to_client(server_set, **client_out_1)
    client_out_2 = psi.client_output(server_out, **client_state)
    return client_out_2

def run_psi_2_round_elgamal(server_set, client_set):
    psi = PSI2RoundElGamal()
    client_out_1, client_state = psi.client_to_server(client_set)
    server_out = psi.server_to_client(server_set, **client_out_1)
    client_out_2 = psi.client_output(server_out, **client_state)
    return client_out_2

if __name__ == '__main__':
    set_len = input('Input set length: ')
    intersection_len = input('Input intersection length: ')

    server_set = []
    client_set = []
    while not (len(client_set) == len(server_set) == set_len):
        server_set = list(set([pyrandom.randint(1, set_len * 10) for i in range(set_len * 5)]))[:set_len]
        client_set = list(set([pyrandom.randint(set_len * 10, set_len * 20) for i in range(set_len * 5)]))[
                     :set_len - intersection_len] + server_set[:intersection_len]

    print
    print('server set: {0}'.format(sorted(server_set)))
    print('client set: {0}'.format(sorted(client_set)))
    print('intersection: {0}'.format(sorted(set(server_set) & set(client_set))))
    print

    tests = [['PSI3RoundPaillier', run_psi_3_round_paillier],
             ['PSI2RoundPaillier', run_psi_2_round_paillier],
             ['PSI2RoundElGamal', run_psi_2_round_elgamal]]

    for test in tests:
        time_taken, result = timer(test[1], server_set, client_set)
        print('{0} output: {1}'.format(test[0], sorted(result)))
        print('{0} time: {1}'.format(test[0], time_taken))
        print<commit_msg>- added comments to simulate_psi.py<commit_after>__author__ = 'Milinda Perera'

import random as pyrandom
from time import time

from psi_3_round_paillier import PSI3RoundPaillier
from psi_2_round_paillier import PSI2RoundPaillier
from psi_2_round_elgamal import PSI2RoundElGamal

def timer(func, *pargs, **kargs):
    """
    Measures the time required to run func with the given parameters.
    Returns the time as well as the result of the computation.
    """
    start = time()
    ret = func(*pargs, **kargs)
    elapsed = time() - start
    return elapsed, ret

def run_psi_3_round_paillier(server_set, client_set):
    """
    Simulates the 3-round PSI protocol based on Paillier encryption scheme
    on the given server and client sets. Returns the final output of the client.
    """
    psi = PSI3RoundPaillier(1024)
    server_out_1, server_state = psi.server_to_client_1(server_set)
    client_out_1 = psi.client_to_server(client_set, **server_out_1)
    server_out_2 = psi.server_to_client_2(client_out_1, **server_state)
    client_out_2 = psi.client_output(server_out_2)
    return client_out_2

def run_psi_2_round_paillier(server_set, client_set):
    """
    Simulates the 2-round PSI protocol based on Paillier encryption scheme
    on the given server and client sets. Returns the final output of the client.
    """
    psi = PSI2RoundPaillier(1024)
    client_out_1, client_state = psi.client_to_server(client_set)
    server_out = psi.server_to_client(server_set, **client_out_1)
    client_out_2 = psi.client_output(server_out, **client_state)
    return client_out_2

def run_psi_2_round_elgamal(server_set, client_set):
    """
    Simulates the 2-round PSI protocol based on ElGamal encryption scheme
    on the given server and client sets. Returns the final output of the client.
    """
    psi = PSI2RoundElGamal()
    client_out_1, client_state = psi.client_to_server(client_set)
    server_out = psi.server_to_client(server_set, **client_out_1)
    client_out_2 = psi.client_output(server_out, **client_state)
    return client_out_2

if __name__ == '__main__':
    # Obtain the server and client set lengths as well as the intersection length from the user.
    set_len = input('Input set length: ')
    intersection_len = input('Input intersection length: ')

    # Generate server and client sets with above parameters.
    server_set = []
    client_set = []
    while not (len(client_set) == len(server_set) == set_len):
        server_set = list(set([pyrandom.randint(1, set_len * 10) for i in range(set_len * 5)]))[:set_len]
        client_set = list(set([pyrandom.randint(set_len * 10, set_len * 20) for i in range(set_len * 5)]))[:set_len - intersection_len] + server_set[:intersection_len]

    # Print generated sets as well as their intersection for comparison purposes.
    print
    print('server set: {0}'.format(sorted(server_set)))
    print('client set: {0}'.format(sorted(client_set)))
    print('intersection: {0}'.format(sorted(set(server_set) & set(client_set))))
    print

    sims = [['PSI3RoundPaillier', run_psi_3_round_paillier],
            ['PSI2RoundPaillier', run_psi_2_round_paillier],
            ['PSI2RoundElGamal', run_psi_2_round_elgamal]]

    # Simulate the protocols and report results.
    for sim in sims:
        time_taken, result = timer(sim[1], server_set, client_set)
        print('{0} output: {1}'.format(sim[0], sorted(result)))
        print('{0} time: {1}'.format(sim[0], time_taken))
        print<|endoftext|>


<commit_before><commit_msg>Made informative prints optional, based on a debug flag (a keyword argument to the MySQL store constructor)<commit_after><|endoftext|>


<commit_before>import sys
import unittest
import time
import click
import traceback
try:
    from ykman.util import BitflagEnum, TRANSPORT
    from ykman.descriptor import get_descriptors
    from click.testing import CliRunner
    from ykman.cli.__main__ import cli
    click.confirm(
        "Run integration tests? This will erase data on the YubiKey,"
        " make sure it is a key used for development.", abort=True)
except Exception:
    sys.exit()

def ykman_cli(*argv):
    runner = CliRunner()
    result = runner.invoke(cli, list(argv), obj={})
    if result.exit_code != 0:
        click.echo(result.output)
        traceback.print_tb(result.exc_info[2])
        raise result.exception
    return result.output

def _one_yubikey():
    return len(list(get_descriptors())) == 1

def _has_mode(mode):
    yubikeys = list(get_descriptors())
    if len(yubikeys) is not 1:
        return False
    return yubikeys[0].mode.has_transport(mode)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
class TestYkmanInfo(unittest.TestCase):

    def test_ykman_info(self):
        time.sleep(3)
        info = ykman_cli('info')
        self.assertIn('Device name:', info)
        self.assertIn('Serial number:', info)
        self.assertIn('Firmware version:', info)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
@unittest.skipIf(not _has_mode(TRANSPORT.OTP), "OTP needs to be enabled")
class TestSlotStatus(unittest.TestCase):

    def test_ykman_slot_info(self):
        info = ykman_cli('slot', 'info')
        self.assertIn('Slot 1:', info)
        self.assertIn('Slot 2:', info)

    def test_ykman_swap_slots(self):
        output = ykman_cli('slot', 'swap', '-f')
        self.assertIn('Swapping slots...', output)
        output = ykman_cli('slot', 'swap', '-f')
        self.assertIn('Swapping slots...', output)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
@unittest.skipIf(not _has_mode(TRANSPORT.OTP), "OTP needs to be enabled")
class TestSlotProgramming(unittest.TestCase):

    def test_ykman_program_otp_slot_2(self):
        output = ykman_cli('slot', 'otp', '2', '-f')
        self.assertIn('Using serial as public ID:', output)
        self.assertIn('Using a randomly generated private ID:', output)
        self.assertIn('Using a randomly generated secret key:', output)
        self._check_slot_2_programmed()

    def test_ykman_program_chalresp_slot_2(self):
        output = ykman_cli('slot', 'chalresp', '2', '-f')
        self.assertIn('Using a randomly generated key.', output)
        self._check_slot_2_programmed()

    def test_ykman_program_hotp_slot_2(self):
        output = ykman_cli(
            'slot', 'hotp', '2',
            '27KIZZE3SD7GE2FVJJBAXEI3I6RRTPGM', '-f')
        self.assertIn('Programming HOTP credential in slot 2...', output)
        self._check_slot_2_programmed()

    def test_ykman_program_static_slot_2(self):
        output = ykman_cli(
            'slot', 'static', '2',
            'higngdukgerjktbbikrhirngtlkkttkb', '-f')
        self.assertIn('Setting static password in slot 2...', output)
        self._check_slot_2_programmed()

    def test_update_settings_enter_slot_2(self):
        ykman_cli('slot', 'otp', '2', '-f')
        output = ykman_cli('slot', 'update', '2', '-f', '--no-enter')
        self.assertIn('Updating settings for slot', output)

    def test_delete_slot_2(self):
        ykman_cli('slot', 'otp', '2', '-f')
        output = ykman_cli('slot', 'delete', '2', '-f')
        self.assertIn('Deleting slot', output)
        status = ykman_cli('slot', 'info')
        self.assertIn('Slot 2: empty', status)

    def _check_slot_2_programmed(self):
        status = ykman_cli('slot', 'info')
        self.assertIn('Slot 2: programmed', status)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
@unittest.skipIf(
    not _has_mode(TRANSPORT.CCID),
    "CCID needs to be enabled for this test.")
class TestOpenPGP(unittest.TestCase):

    def test_openpgp_info(self):
        output = ykman_cli('openpgp', 'info')
        self.assertIn('OpenPGP version:', output)

    def test_openpgp_reset(self):
        output = ykman_cli('openpgp', 'reset', '-f')
        self.assertIn(
            'Success! All data has been cleared and default PINs are set.',
            output)
<commit_msg>Fix flake8<commit_after>import sys
import unittest
import time
import click
import traceback
try:
    from ykman.util import TRANSPORT
    from ykman.descriptor import get_descriptors
    from click.testing import CliRunner
    from ykman.cli.__main__ import cli
    click.confirm(
        "Run integration tests? This will erase data on the YubiKey,"
        " make sure it is a key used for development.", abort=True)
except Exception:
    sys.exit()

def ykman_cli(*argv):
    runner = CliRunner()
    result = runner.invoke(cli, list(argv), obj={})
    if result.exit_code != 0:
        click.echo(result.output)
        traceback.print_tb(result.exc_info[2])
        raise result.exception
    return result.output

def _one_yubikey():
    return len(list(get_descriptors())) == 1

def _has_mode(mode):
    yubikeys = list(get_descriptors())
    if len(yubikeys) is not 1:
        return False
    return yubikeys[0].mode.has_transport(mode)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
class TestYkmanInfo(unittest.TestCase):

    def test_ykman_info(self):
        time.sleep(3)
        info = ykman_cli('info')
        self.assertIn('Device name:', info)
        self.assertIn('Serial number:', info)
        self.assertIn('Firmware version:', info)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
@unittest.skipIf(not _has_mode(TRANSPORT.OTP), "OTP needs to be enabled")
class TestSlotStatus(unittest.TestCase):

    def test_ykman_slot_info(self):
        info = ykman_cli('slot', 'info')
        self.assertIn('Slot 1:', info)
        self.assertIn('Slot 2:', info)

    def test_ykman_swap_slots(self):
        output = ykman_cli('slot', 'swap', '-f')
        self.assertIn('Swapping slots...', output)
        output = ykman_cli('slot', 'swap', '-f')
        self.assertIn('Swapping slots...', output)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
@unittest.skipIf(not _has_mode(TRANSPORT.OTP), "OTP needs to be enabled")
class TestSlotProgramming(unittest.TestCase):

    def test_ykman_program_otp_slot_2(self):
        output = ykman_cli('slot', 'otp', '2', '-f')
        self.assertIn('Using serial as public ID:', output)
        self.assertIn('Using a randomly generated private ID:', output)
        self.assertIn('Using a randomly generated secret key:', output)
        self._check_slot_2_programmed()

    def test_ykman_program_chalresp_slot_2(self):
        output = ykman_cli('slot', 'chalresp', '2', '-f')
        self.assertIn('Using a randomly generated key.', output)
        self._check_slot_2_programmed()

    def test_ykman_program_hotp_slot_2(self):
        output = ykman_cli(
            'slot', 'hotp', '2',
            '27KIZZE3SD7GE2FVJJBAXEI3I6RRTPGM', '-f')
        self.assertIn('Programming HOTP credential in slot 2...', output)
        self._check_slot_2_programmed()

    def test_ykman_program_static_slot_2(self):
        output = ykman_cli(
            'slot', 'static', '2',
            'higngdukgerjktbbikrhirngtlkkttkb', '-f')
        self.assertIn('Setting static password in slot 2...', output)
        self._check_slot_2_programmed()

    def test_update_settings_enter_slot_2(self):
        ykman_cli('slot', 'otp', '2', '-f')
        output = ykman_cli('slot', 'update', '2', '-f', '--no-enter')
        self.assertIn('Updating settings for slot', output)

    def test_delete_slot_2(self):
        ykman_cli('slot', 'otp', '2', '-f')
        output = ykman_cli('slot', 'delete', '2', '-f')
        self.assertIn('Deleting slot', output)
        status = ykman_cli('slot', 'info')
        self.assertIn('Slot 2: empty', status)

    def _check_slot_2_programmed(self):
        status = ykman_cli('slot', 'info')
        self.assertIn('Slot 2: programmed', status)

@unittest.skipIf(not _one_yubikey(), "A single YubiKey need to be connected.")
@unittest.skipIf(
    not _has_mode(TRANSPORT.CCID),
    "CCID needs to be enabled for this test.")
class TestOpenPGP(unittest.TestCase):

    def test_openpgp_info(self):
        output = ykman_cli('openpgp', 'info')
        self.assertIn('OpenPGP version:', output)

    def test_openpgp_reset(self):
        output = ykman_cli('openpgp', 'reset', '-f')
        self.assertIn(
            'Success! All data has been cleared and default PINs are set.',
            output)
<|endoftext|>