<commit_before># Copyright (C) 2011 Nippon Telegraph and Telephone Corporation.
# OpenFlow 1.2 upgrade, modified for peering exchanges by Sam Russell 2013
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import struct

from ryu.base import app_manager
from ryu.controller import mac_to_port
from ryu.controller import ofp_event
from ryu.topology import event
from ryu.controller.handler import MAIN_DISPATCHER, DEAD_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.controller.handler import set_ev_handler
from ryu.ofproto import ofproto_v1_2, ofproto_v1_2_parser, ether
from ryu.lib.mac import haddr_to_str
from ryu.lib import ip
from ryu.lib.packet import packet
import spe_config

# TODO: we should split the handler into two parts, protocol
# independent and dependant parts.

# TODO: can we use dpkt python library?

# TODO: we need to move the followings to something like db

class SPE(app_manager.RyuApp):
    OFP_VERSIONS = [ofproto_v1_2.OFP_VERSION]

    def __init__(self, *args, **kwargs):
        super(SPE, self).__init__(*args, **kwargs)
        self.dps = {}
        self.mac_to_port = {}

    def add_flow(self, datapath, table_id, match, instructions, priority=0x8000, buffer_id=ofproto_v1_2.OFP_NO_BUFFER):
        ofproto = datapath.ofproto

        #wildcards = ofproto_v1_2.OFPFW_ALL
        #wildcards &= ~ofproto_v1_2.OFPFW_IN_PORT
        #wildcards &= ~ofproto_v1_2.OFPFW_DL_DST
        # old standard matching - deprecated from 1.2 onwards

        #match = datapath.ofproto_parser.OFPMatch(
        #    wildcards, in_port, 0, dst,
        #    0, 0, 0, 0, 0, 0, 0, 0, 0)

        #datapath, cookie, cookie_mask, table_id, command,
        #         idle_timeout, hard_timeout, priority, buffer_id, out_port,
        #         out_group, flags, match, instructions)
        mod = datapath.ofproto_parser.OFPFlowMod(
            datapath=datapath, cookie=0, cookie_mask=0, table_id=table_id,
            command=ofproto.OFPFC_ADD, idle_timeout=0, hard_timeout=0,
            priority=0x8000, buffer_id=buffer_id,
            out_port=ofproto_v1_2.OFPP_ANY, out_group=ofproto_v1_2.OFPG_ANY,
            flags=ofproto.OFPFF_SEND_FLOW_REM, match=match, instructions=instructions)
        datapath.send_msg(mod)

    def add_arp_reply_catcher(self, datapath, ipaddr = None, port = None, table_id=0):
        match = datapath.ofproto_parser.OFPMatch()
        match.set_dl_type(ether.ETH_TYPE_ARP)
        match.set_arp_opcode(2)
        if port:
            match.set_in_port(port)
        if ipaddr:
            match.set_arp_spa(ipaddr)
        ofproto = datapath.ofproto
        instructions = [datapath.ofproto_parser.OFPInstructionGotoTable(2)]
        self.add_flow(datapath, table_id, match, instructions, priority=0x8000)

    def add_arp_request_forwarder(self, datapath, ipaddr, port, table_id=0):
        match = datapath.ofproto_parser.OFPMatch()
        match.set_dl_type(ether.ETH_TYPE_ARP)
        match.set_arp_opcode(1)
        match.set_arp_tpa(ipaddr)
        ofproto = datapath.ofproto
        actions = [datapath.ofproto_parser.OFPActionOutput(port, 1500)]
        instructions = [datapath.ofproto_parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS, actions)]
        self.add_flow(datapath, table_id, match, instructions, priority=0x8000)

    def init_flows(self, datapath):
        ofproto = datapath.ofproto
        # send arp replies to controller always
        self.add_arp_reply_catcher(datapath)
        for port, ipaddr in spe_config.ports.iteritems():
            self.add_arp_reply_catcher(datapath, ipaddr=ip.ipv4_to_bin(ipaddr), port=port, table_id=1)
            self.add_arp_request_forwarder(datapath, ipaddr=ip.ipv4_to_bin(ipaddr), port=port, table_id=1)

        # make a flow to send all ARP packets to table 1
        match = datapath.ofproto_parser.OFPMatch()
        match.set_dl_type(ether.ETH_TYPE_ARP)
        instructions = [datapath.ofproto_parser.OFPInstructionGotoTable(1)]
        self.add_flow(datapath, 0, match, instructions)

    @set_ev_cls(ofp_event.EventOFPStateChange, [MAIN_DISPATCHER, DEAD_DISPATCHER])
    def switch_enter_handler(self, ev):
        dp = ev.datapath
        if ev.state == MAIN_DISPATCHER:
            self.logger.info("Switch entered: %s", dp.id)
            self.init_flows(dp)

        elif ev.state == DEAD_DISPATCHER:
            if dp.id is None:
                return
            self.logger.info("Switch left: %s", dp.id)

    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def _packet_in_handler(self, ev):
        msg = ev.msg
        datapath = msg.datapath
        ofproto = datapath.ofproto

        dst, src, _eth_type = struct.unpack_from('!6s6sH', buffer(msg.data), 0)

        dpid = datapath.id
        self.mac_to_port.setdefault(dpid, {})

        match = msg.match
        in_port = 0
        ethtype = 0
        #iterate through fields - parser should handle this
        #packet in dpid 20015998343868 from 08:00:27:15:d4:53 to ff:ff:ff:ff:ff:ff log_port 0 phy_port 0
        #Field MTInPort(header=2147483652,length=8,n_bytes=4,value=2)
        #Field MTEthType(header=2147486210,length=6,n_bytes=2,value=2054)
        #Field MTArpOp(header=2147494402,length=6,n_bytes=2,value=1)
        #Field MTMetadata(header=2147484680,length=12,n_bytes=8,value=18446744073709551615L)
        #Field MTArpSha(header=2147495942,length=10,n_bytes=6,value="\x08\x00'\x15\xd4S")
        #Field MTEthDst(header=2147485190,length=10,n_bytes=6,value='\xff\xff\xff\xff\xff\xff')
        #Field MTArpSpa(header=2147494916,length=8,n_bytes=4,value=167772161)
        #Field MTArpTha(header=2147496454,length=10,n_bytes=6,value='\x00\x00\x00\x00\x00\x00')

        # we should build a dictionary of MTXXXX class names to speed this up
        for o in match.fields:
            self.logger.info("Field %s", str(o))
            if isinstance(o, ofproto_v1_2_parser.MTInPort):
                in_port = o.value
                break

        self.logger.info("packet in dpid %s from %s to %s log_port %s",
                         dpid, haddr_to_str(src), haddr_to_str(dst),
                         in_port)

        # parse packet

        pkt = packet.Packet(msg.data)
        eth_pkt = pkt.next()
        ethtype = eth_pkt.ethertype

        # if ARP (request) then flood and don't make a flow
        if ethtype == ether.ETH_TYPE_ARP:
            # if ARP reply then drop
            arp_pkt = pkt.next()
            if arp_pkt.opcode == 2:
                return
            out_port = ofproto_v1_2.OFPP_FLOOD
            actions = [datapath.ofproto_parser.OFPActionOutput(out_port, 1500)]
            out = datapath.ofproto_parser.OFPPacketOut(
                datapath=datapath, buffer_id=msg.buffer_id, in_port=in_port,
                actions=actions)
            datapath.send_msg(out)

        # do we know the mac?
        if src not in self.mac_to_port[dpid]:
            # learn the mac address to avoid FLOOD next time.
            self.mac_to_port[dpid][src] = in_port
            # set a flow to table 0 to allow packets through to table 1
            match = datapath.ofproto_parser.OFPMatch()
            match.set_in_port(in_port)
            match.set_dl_src(src)
            match.set_dl_type(ethtype)
            instructions = [datapath.ofproto_parser.OFPInstructionGotoTable(2)]
            self.add_flow(datapath, 0, match, instructions)

        if dst in self.mac_to_port[dpid]:
            out_port = self.mac_to_port[dpid][dst]
            match = datapath.ofproto_parser.OFPMatch()
            match.set_dl_dst(dst)
            actions = [datapath.ofproto_parser.OFPActionOutput(out_port, 1500)]
            instructions = [datapath.ofproto_parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS, actions)]
            self.add_flow(datapath, 2, match, instructions, buffer_id=msg.buffer_id)
        else:
            out_port = ofproto_v1_2.OFPP_FLOOD
            actions = [datapath.ofproto_parser.OFPActionOutput(out_port, 1500)]
            out = datapath.ofproto_parser.OFPPacketOut(
                datapath=datapath, buffer_id=msg.buffer_id, in_port=in_port,
                actions=actions)
            datapath.send_msg(out)

    @set_ev_cls(ofp_event.EventOFPPortStatus, MAIN_DISPATCHER)
    def _port_status_handler(self, ev):
        msg = ev.msg
        reason = msg.reason
        port_no = msg.desc.port_no

        ofproto = msg.datapath.ofproto
        if reason == ofproto.OFPPR_ADD:
            self.logger.info("port added %s", port_no)
        elif reason == ofproto.OFPPR_DELETE:
            self.logger.info("port deleted %s", port_no)
        elif reason == ofproto.OFPPR_MODIFY:
            self.logger.info("port modified %s", port_no)
        else:
            self.logger.info("Illeagal port state %s %s", port_no, reason)
<commit_msg>got rid of arpreply to table 1<commit_after># Copyright (C) 2011 Nippon Telegraph and Telephone Corporation.
# OpenFlow 1.2 upgrade, modified for peering exchanges by Sam Russell 2013
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import struct

from ryu.base import app_manager
from ryu.controller import mac_to_port
from ryu.controller import ofp_event
from ryu.topology import event
from ryu.controller.handler import MAIN_DISPATCHER, DEAD_DISPATCHER
from ryu.controller.handler import set_ev_cls
from ryu.controller.handler import set_ev_handler
from ryu.ofproto import ofproto_v1_2, ofproto_v1_2_parser, ether
from ryu.lib.mac import haddr_to_str
from ryu.lib import ip
from ryu.lib.packet import packet
import spe_config

# TODO: we should split the handler into two parts, protocol
# independent and dependant parts.

# TODO: can we use dpkt python library?

# TODO: we need to move the followings to something like db

class SPE(app_manager.RyuApp):
    OFP_VERSIONS = [ofproto_v1_2.OFP_VERSION]

    def __init__(self, *args, **kwargs):
        super(SPE, self).__init__(*args, **kwargs)
        self.dps = {}
        self.mac_to_port = {}

    def add_flow(self, datapath, table_id, match, instructions, priority=0x8000, buffer_id=ofproto_v1_2.OFP_NO_BUFFER):
        ofproto = datapath.ofproto

        #wildcards = ofproto_v1_2.OFPFW_ALL
        #wildcards &= ~ofproto_v1_2.OFPFW_IN_PORT
        #wildcards &= ~ofproto_v1_2.OFPFW_DL_DST
        # old standard matching - deprecated from 1.2 onwards

        #match = datapath.ofproto_parser.OFPMatch(
        #    wildcards, in_port, 0, dst,
        #    0, 0, 0, 0, 0, 0, 0, 0, 0)

        #datapath, cookie, cookie_mask, table_id, command,
        #         idle_timeout, hard_timeout, priority, buffer_id, out_port,
        #         out_group, flags, match, instructions)
        mod = datapath.ofproto_parser.OFPFlowMod(
            datapath=datapath, cookie=0, cookie_mask=0, table_id=table_id,
            command=ofproto.OFPFC_ADD, idle_timeout=0, hard_timeout=0,
            priority=0x8000, buffer_id=buffer_id,
            out_port=ofproto_v1_2.OFPP_ANY, out_group=ofproto_v1_2.OFPG_ANY,
            flags=ofproto.OFPFF_SEND_FLOW_REM, match=match, instructions=instructions)
        datapath.send_msg(mod)

    def add_arp_reply_catcher(self, datapath, ipaddr, port, table_id=0):
        match = datapath.ofproto_parser.OFPMatch()
        match.set_dl_type(ether.ETH_TYPE_ARP)
        match.set_arp_opcode(2)
        if port:
            match.set_in_port(port)
        if ipaddr:
            match.set_arp_spa(ipaddr)
        ofproto = datapath.ofproto
        instructions = [datapath.ofproto_parser.OFPInstructionGotoTable(2)]
        self.add_flow(datapath, table_id, match, instructions, priority=0x8000)

    def add_arp_request_forwarder(self, datapath, ipaddr, port, table_id=0):
        match = datapath.ofproto_parser.OFPMatch()
        match.set_dl_type(ether.ETH_TYPE_ARP)
        match.set_arp_opcode(1)
        match.set_arp_tpa(ipaddr)
        ofproto = datapath.ofproto
        actions = [datapath.ofproto_parser.OFPActionOutput(port, 1500)]
        instructions = [datapath.ofproto_parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS, actions)]
        self.add_flow(datapath, table_id, match, instructions, priority=0x8000)

    def init_flows(self, datapath):
        ofproto = datapath.ofproto
        for port, ipaddr in spe_config.ports.iteritems():
            self.add_arp_reply_catcher(datapath, ipaddr=ip.ipv4_to_bin(ipaddr), port=port, table_id=1)
            self.add_arp_request_forwarder(datapath, ipaddr=ip.ipv4_to_bin(ipaddr), port=port, table_id=1)

        # make a flow to send all ARP packets to table 1
        match = datapath.ofproto_parser.OFPMatch()
        match.set_dl_type(ether.ETH_TYPE_ARP)
        instructions = [datapath.ofproto_parser.OFPInstructionGotoTable(1)]
        self.add_flow(datapath, 0, match, instructions)

    @set_ev_cls(ofp_event.EventOFPStateChange, [MAIN_DISPATCHER, DEAD_DISPATCHER])
    def switch_enter_handler(self, ev):
        dp = ev.datapath
        if ev.state == MAIN_DISPATCHER:
            self.logger.info("Switch entered: %s", dp.id)
            self.init_flows(dp)

        elif ev.state == DEAD_DISPATCHER:
            if dp.id is None:
                return
            self.logger.info("Switch left: %s", dp.id)

    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def _packet_in_handler(self, ev):
        msg = ev.msg
        datapath = msg.datapath
        ofproto = datapath.ofproto

        dst, src, _eth_type = struct.unpack_from('!6s6sH', buffer(msg.data), 0)

        dpid = datapath.id
        self.mac_to_port.setdefault(dpid, {})

        match = msg.match
        in_port = 0
        ethtype = 0
        #iterate through fields - parser should handle this
        #packet in dpid 20015998343868 from 08:00:27:15:d4:53 to ff:ff:ff:ff:ff:ff log_port 0 phy_port 0
        #Field MTInPort(header=2147483652,length=8,n_bytes=4,value=2)
        #Field MTEthType(header=2147486210,length=6,n_bytes=2,value=2054)
        #Field MTArpOp(header=2147494402,length=6,n_bytes=2,value=1)
        #Field MTMetadata(header=2147484680,length=12,n_bytes=8,value=18446744073709551615L)
        #Field MTArpSha(header=2147495942,length=10,n_bytes=6,value="\x08\x00'\x15\xd4S")
        #Field MTEthDst(header=2147485190,length=10,n_bytes=6,value='\xff\xff\xff\xff\xff\xff')
        #Field MTArpSpa(header=2147494916,length=8,n_bytes=4,value=167772161)
        #Field MTArpTha(header=2147496454,length=10,n_bytes=6,value='\x00\x00\x00\x00\x00\x00')

        # we should build a dictionary of MTXXXX class names to speed this up
        for o in match.fields:
            self.logger.info("Field %s", str(o))
            if isinstance(o, ofproto_v1_2_parser.MTInPort):
                in_port = o.value
                break

        self.logger.info("packet in dpid %s from %s to %s log_port %s",
                         dpid, haddr_to_str(src), haddr_to_str(dst),
                         in_port)

        # parse packet

        pkt = packet.Packet(msg.data)
        eth_pkt = pkt.next()
        ethtype = eth_pkt.ethertype

        # if ARP (request) then flood and don't make a flow
        if ethtype == ether.ETH_TYPE_ARP:
            # if ARP reply then drop
            arp_pkt = pkt.next()
            if arp_pkt.opcode == 2:
                return
            out_port = ofproto_v1_2.OFPP_FLOOD
            actions = [datapath.ofproto_parser.OFPActionOutput(out_port, 1500)]
            out = datapath.ofproto_parser.OFPPacketOut(
                datapath=datapath, buffer_id=msg.buffer_id, in_port=in_port,
                actions=actions)
            datapath.send_msg(out)

        # do we know the mac?
        if src not in self.mac_to_port[dpid]:
            # learn the mac address to avoid FLOOD next time.
            self.mac_to_port[dpid][src] = in_port
            # set a flow to table 0 to allow packets through to table 1
            match = datapath.ofproto_parser.OFPMatch()
            match.set_in_port(in_port)
            match.set_dl_src(src)
            match.set_dl_type(ethtype)
            instructions = [datapath.ofproto_parser.OFPInstructionGotoTable(2)]
            self.add_flow(datapath, 0, match, instructions)

        if dst in self.mac_to_port[dpid]:
            out_port = self.mac_to_port[dpid][dst]
            match = datapath.ofproto_parser.OFPMatch()
            match.set_dl_dst(dst)
            actions = [datapath.ofproto_parser.OFPActionOutput(out_port, 1500)]
            instructions = [datapath.ofproto_parser.OFPInstructionActions(ofproto.OFPIT_APPLY_ACTIONS, actions)]
            self.add_flow(datapath, 2, match, instructions, buffer_id=msg.buffer_id)
        else:
            out_port = ofproto_v1_2.OFPP_FLOOD
            actions = [datapath.ofproto_parser.OFPActionOutput(out_port, 1500)]
            out = datapath.ofproto_parser.OFPPacketOut(
                datapath=datapath, buffer_id=msg.buffer_id, in_port=in_port,
                actions=actions)
            datapath.send_msg(out)

    @set_ev_cls(ofp_event.EventOFPPortStatus, MAIN_DISPATCHER)
    def _port_status_handler(self, ev):
        msg = ev.msg
        reason = msg.reason
        port_no = msg.desc.port_no

        ofproto = msg.datapath.ofproto
        if reason == ofproto.OFPPR_ADD:
            self.logger.info("port added %s", port_no)
        elif reason == ofproto.OFPPR_DELETE:
            self.logger.info("port deleted %s", port_no)
        elif reason == ofproto.OFPPR_MODIFY:
            self.logger.info("port modified %s", port_no)
        else:
            self.logger.info("Illeagal port state %s %s", port_no, reason)
<|endoftext|>


<commit_before>"""Support for esphome sensors."""
import math
from typing import Optional

from aioesphomeapi import SensorInfo, SensorState, TextSensorInfo, TextSensorState

from homeassistant.config_entries import ConfigEntry
from homeassistant.helpers.typing import HomeAssistantType

from . import EsphomeEntity, esphome_state_property, platform_async_setup_entry

async def async_setup_entry(
    hass: HomeAssistantType, entry: ConfigEntry, async_add_entities
) -> None:
    """Set up esphome sensors based on a config entry."""
    await platform_async_setup_entry(
        hass,
        entry,
        async_add_entities,
        component_key="sensor",
        info_type=SensorInfo,
        entity_type=EsphomeSensor,
        state_type=SensorState,
    )
    await platform_async_setup_entry(
        hass,
        entry,
        async_add_entities,
        component_key="text_sensor",
        info_type=TextSensorInfo,
        entity_type=EsphomeTextSensor,
        state_type=TextSensorState,
    )

# https://github.com/PyCQA/pylint/issues/3150 for all @esphome_state_property
# pylint: disable=invalid-overridden-method

class EsphomeSensor(EsphomeEntity):
    """A sensor implementation for esphome."""

    @property
    def _static_info(self) -> SensorInfo:
        return super()._static_info

    @property
    def _state(self) -> Optional[SensorState]:
        return super()._state

    @property
    def icon(self) -> str:
        """Return the icon."""
        if self._static_info.icon == "":
            return None
        return self._static_info.icon

    @property
    def force_update(self) -> bool:
        """Return if this sensor should force a state update."""
        return self._static_info.force_update

    @esphome_state_property
    def state(self) -> Optional[str]:
        """Return the state of the entity."""
        if math.isnan(self._state.state):
            return None
        if self._state.missing_state:
            return None
        return f"{self._state.state:.{self._static_info.accuracy_decimals}f}"

    @property
    def unit_of_measurement(self) -> str:
        """Return the unit the value is expressed in."""
        if self._static_info.unit_of_measurement == "":
            return None
        return self._static_info.unit_of_measurement

    @property
    def device_class(self) -> str:
        """Return the class of this device, from component DEVICE_CLASSES."""
        if self._static_info.device_class == "":
            return None
        return self._static_info.device_class

class EsphomeTextSensor(EsphomeEntity):
    """A text sensor implementation for ESPHome."""

    @property
    def _static_info(self) -> TextSensorInfo:
        return super()._static_info

    @property
    def _state(self) -> Optional[TextSensorState]:
        return super()._state

    @property
    def icon(self) -> str:
        """Return the icon."""
        return self._static_info.icon

    @esphome_state_property
    def state(self) -> Optional[str]:
        """Return the state of the entity."""
        if self._state.missing_state:
            return None
        return self._state.state
<commit_msg>Do not provide icon if device class is set in ESPHome config (#46650)<commit_after>"""Support for esphome sensors."""
import math
from typing import Optional

from aioesphomeapi import SensorInfo, SensorState, TextSensorInfo, TextSensorState

from homeassistant.config_entries import ConfigEntry
from homeassistant.helpers.typing import HomeAssistantType

from . import EsphomeEntity, esphome_state_property, platform_async_setup_entry

async def async_setup_entry(
    hass: HomeAssistantType, entry: ConfigEntry, async_add_entities
) -> None:
    """Set up esphome sensors based on a config entry."""
    await platform_async_setup_entry(
        hass,
        entry,
        async_add_entities,
        component_key="sensor",
        info_type=SensorInfo,
        entity_type=EsphomeSensor,
        state_type=SensorState,
    )
    await platform_async_setup_entry(
        hass,
        entry,
        async_add_entities,
        component_key="text_sensor",
        info_type=TextSensorInfo,
        entity_type=EsphomeTextSensor,
        state_type=TextSensorState,
    )

# https://github.com/PyCQA/pylint/issues/3150 for all @esphome_state_property
# pylint: disable=invalid-overridden-method

class EsphomeSensor(EsphomeEntity):
    """A sensor implementation for esphome."""

    @property
    def _static_info(self) -> SensorInfo:
        return super()._static_info

    @property
    def _state(self) -> Optional[SensorState]:
        return super()._state

    @property
    def icon(self) -> str:
        """Return the icon."""
        if not self._static_info.icon or self._static_info.device_class:
            return None
        return self._static_info.icon

    @property
    def force_update(self) -> bool:
        """Return if this sensor should force a state update."""
        return self._static_info.force_update

    @esphome_state_property
    def state(self) -> Optional[str]:
        """Return the state of the entity."""
        if math.isnan(self._state.state):
            return None
        if self._state.missing_state:
            return None
        return f"{self._state.state:.{self._static_info.accuracy_decimals}f}"

    @property
    def unit_of_measurement(self) -> str:
        """Return the unit the value is expressed in."""
        if not self._static_info.unit_of_measurement:
            return None
        return self._static_info.unit_of_measurement

    @property
    def device_class(self) -> str:
        """Return the class of this device, from component DEVICE_CLASSES."""
        if not self._static_info.device_class:
            return None
        return self._static_info.device_class

class EsphomeTextSensor(EsphomeEntity):
    """A text sensor implementation for ESPHome."""

    @property
    def _static_info(self) -> TextSensorInfo:
        return super()._static_info

    @property
    def _state(self) -> Optional[TextSensorState]:
        return super()._state

    @property
    def icon(self) -> str:
        """Return the icon."""
        return self._static_info.icon

    @esphome_state_property
    def state(self) -> Optional[str]:
        """Return the state of the entity."""
        if self._state.missing_state:
            return None
        return self._state.state
<|endoftext|>


<commit_before>import logging
from django.core.mail import send_mail
from django.core.management.base import  BaseCommand
from casexml.apps.case.models import CommCareCase
from corehq.pillows import CasePillow

CHUNK_SIZE=500
POOL_SIZE = 15

class Command(BaseCommand):
    help = "Fast reindex of case elastic index by using the case view and reindexing cases"

    def handle(self, *args, **options):
        confirm = raw_input("""
    You have requested to do a case elastic index reset via fast track.
    This will IRREVERSIBLY REMOVE
    ALL index data in the case index and will take a while to reload.
    Are you sure you want to do this. Also you MUST have run_ptop disabled for this to run.

    Type 'yes' to continue, or 'no' to cancel: """)

        if confirm != 'yes':
            print "Reset cancelled."
            return

        confirm_ptop = raw_input("""Are you sure you disabled run_ptop? """)
        if confirm_ptop != "yes":
            return

        confirm_alias = raw_input("""Are you sure you are not blowing away a production index?""")
        if confirm_alias != "yes":
            return

        #delete the existing index.
        casepillow = CasePillow()
        print "Deleting index"
        casepillow.delete_index()
        print "Recreating index"
        casepillow.create_index()
        print "Resetting CasePillow Checkpoint"

        casepillow.reset_checkpoint()

        db = CommCareCase.get_db()
        start_num = 0

        print "starting fast tracked reindexing"
        chunk = db.view('case/by_owner', reduce=False, limit=CHUNK_SIZE, skip=start_num)

        #though this might cause some superfluous reindexes of docs,
        # we're going to set the checkpoint BEFORE we start our operation so that any changes
        # that happen to cases while we're doing our reindexing would not get skipped once we
        # finish.
        casepillow.set_checkpoint({ 'seq': casepillow.couch_db.info()['update_seq'] } )
        def do_index(item):
            print "Processing: %s" % item['id']
            casepillow.processor(item, do_set_checkpoint=False)

        try:
            while len(chunk) > 0:
                for item in chunk:
                    casepillow.processor(item, do_set_checkpoint=False)
                start_num += CHUNK_SIZE
                print "Grabbing next chunk: %d" % start_num
                chunk = db.view('case/by_owner', reduce=False, limit=CHUNK_SIZE, skip=start_num)

            print "Index recreated - you may now restart run_ptop"
            send_mail('[commcare-hq] Pillowtop Case Reindex Complete',
                "Case reindex complete for index %s - it may now be aliased to hqcases" % casepillow
                .es_index,
                'hq-noreply@dimagi.com', ['commcarehq-dev@dimagi.com'], fail_silently=True)
        except Exception, ex:
            logging.exception("Case pillowtop fast reindex failed: %s" % ex)

<commit_msg>clear the type cache when reindexing the cases fast track<commit_after>import logging
from django.core.mail import send_mail
from django.core.management.base import  BaseCommand
from casexml.apps.case.models import CommCareCase
from corehq.pillows import CasePillow

CHUNK_SIZE=500
POOL_SIZE = 15

class Command(BaseCommand):
    help = "Fast reindex of case elastic index by using the case view and reindexing cases"

    def handle(self, *args, **options):
        confirm = raw_input("""
    You have requested to do a case elastic index reset via fast track.
    This will IRREVERSIBLY REMOVE
    ALL index data in the case index and will take a while to reload.
    Are you sure you want to do this. Also you MUST have run_ptop disabled for this to run.

    Type 'yes' to continue, or 'no' to cancel: """)

        if confirm != 'yes':
            print "Reset cancelled."
            return

        confirm_ptop = raw_input("""Are you sure you disabled run_ptop? """)
        if confirm_ptop != "yes":
            return

        confirm_alias = raw_input("""Are you sure you are not blowing away a production index?""")
        if confirm_alias != "yes":
            return

        #delete the existing index.
        casepillow = CasePillow()
        print "Deleting index"
        casepillow.delete_index()
        print "Recreating index"
        casepillow.create_index()
        casepillow.seen_types = {}
        print "Resetting CasePillow Checkpoint"

        casepillow.reset_checkpoint()

        db = CommCareCase.get_db()
        start_num = 0

        print "starting fast tracked reindexing"
        chunk = db.view('case/by_owner', reduce=False, limit=CHUNK_SIZE, skip=start_num)

        #though this might cause some superfluous reindexes of docs,
        # we're going to set the checkpoint BEFORE we start our operation so that any changes
        # that happen to cases while we're doing our reindexing would not get skipped once we
        # finish.
        casepillow.set_checkpoint({ 'seq': casepillow.couch_db.info()['update_seq'] } )
        def do_index(item):
            print "Processing: %s" % item['id']
            casepillow.processor(item, do_set_checkpoint=False)

        try:
            while len(chunk) > 0:
                for item in chunk:
                    casepillow.processor(item, do_set_checkpoint=False)
                start_num += CHUNK_SIZE
                print "Grabbing next chunk: %d" % start_num
                chunk = db.view('case/by_owner', reduce=False, limit=CHUNK_SIZE, skip=start_num)

            print "Index recreated - you may now restart run_ptop"
            send_mail('[commcare-hq] Pillowtop Case Reindex Complete',
                "Case reindex complete for index %s - it may now be aliased to hqcases" % casepillow
                .es_index,
                'hq-noreply@dimagi.com', ['commcarehq-dev@dimagi.com'], fail_silently=True)
        except Exception, ex:
            logging.exception("Case pillowtop fast reindex failed: %s" % ex)

<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from numpy.testing import assert_approx_equal, assert_allclose
from nose.tools import assert_equal
from UliEngineering.Electronics.TemperatureCoefficient import *
from UliEngineering.EngineerIO import auto_format
import numpy as np

class TestTemperatureCoefficient(object):
    def test_value_range_over_temperature1(self):
        # Test with simple ppm input
        assert_equal(value_range_over_temperature("1 kΩ", "100 ppm"),
            ValueRange("994 Ω", "1.006 kΩ")
        )

    def test_value_range_over_temperature2(self):
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", ("-100 ppm", "100 ppm")),
            ValueRange("994 Ω", "1.006 kΩ")
        )
        # Test with ++ the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", ("100 ppm", "100 ppm")),
            ValueRange("1.006 Ω", "1.006 kΩ")
        )

    def test_value_range_over_temperature3(self):
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", ("0 ppm", "100 ppm")),
            ValueRange("1000 Ω", "1.006 kΩ")
        )

    def test_value_range_over_temperature_percent(self):
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", "1 %"),
            ValueRange("350 Ω", "1.650 kΩ")
        )
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", "1.006 %"),
            ValueRange("346 Ω", "1.654 kΩ")
        )

<commit_msg>Fix tempco unit tests<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from numpy.testing import assert_approx_equal, assert_allclose
from nose.tools import assert_equal
from UliEngineering.Electronics.TemperatureCoefficient import *
from UliEngineering.EngineerIO import auto_format
import numpy as np

class TestTemperatureCoefficient(object):
    def test_value_range_over_temperature1(self):
        # Test with simple ppm input
        assert_equal(value_range_over_temperature("1 kΩ", "100 ppm"),
            ValueRange("994 Ω", "1.006 kΩ")
        )

    def test_value_range_over_temperature2(self):
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", ("-100 ppm", "100 ppm")),
            ValueRange("994 Ω", "1.006 kΩ")
        )
        # Test with ++ the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", ("+100 ppm", "+100 ppm")),
            ValueRange("1.006 kΩ", "1.006 kΩ")
        )

    def test_value_range_over_temperature3(self):
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", ("0 ppm", "100 ppm")),
            ValueRange("1000 Ω", "1.006 kΩ")
        )

    def test_value_range_over_temperature_percent(self):
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", "1 %"),
            ValueRange("350 Ω", "1.650 kΩ")
        )
        # Test with +- the same ppm input
        assert_equal(value_range_over_temperature("1 kΩ", "1.006 %"),
            ValueRange("346 Ω", "1.654 kΩ")
        )

<|endoftext|>


<commit_before>###########
# Parsers #
###########
import decimal
import re
from datetime import datetime
from sqlalchemy import Table, MetaData, Column, select, column, func, literal_column, and_, or_
from api.error import APIError, APIKeyError
from api.connection import _get_engine
import geoalchemy2  # Although this import seems unused is has to be here

pgsql_qualifier = re.compile(r"^[\w\d_\.]+$")

def get_or_403(dictionary, key):
    try:
        return dictionary[key]
    except KeyError:
        raise APIKeyError(dictionary, key)

def is_pg_qual(x):
    return pgsql_qualifier.search(x)

def quote(x):
    if not x.startswith('"') and '(' not in x:
        return '"' + x + '"'
    else:
        return x

def read_pgvalue(x):
    # TODO: Implement check for valid values
    if isinstance(x, str):
        return "'" + x + "'"
    if x is None:
        return 'null'
    return x

def read_operator(x, right):
    # TODO: Implement check for valid operators
    if isinstance(right, dict) and get_or_403(right, 'type') == 'value' and ('value' not in right or get_or_403(right, 'value') is None and x == '='):
        return 'is'
    return x

class ValidationError(Exception):
    def __init__(self, message, value):
        self.message = message
        self.value = value

def read_bool(s):
    if isinstance(s, bool):
        return s
    if s.lower() in ["true", "false"]:
        return s.lower() == "true"
    elif s.lower() in ["yes", "no"]:
        return s.lower() == "true"
    else:
        raise APIError("Invalid value in binary field", s)

def read_pgid(s):
    if is_pg_qual(s):
        return s
    raise APIError("Invalid identifier: '%s'"%s)

def set_meta_info(method, user, message=None):
    val_dict = {}
    val_dict['_user'] = user  # TODO: Add user handling
    val_dict['_message'] = message
    return val_dict

def parse_insert(d, context, message=None):
    table = Table(read_pgid(get_or_403(d, 'table')), MetaData(bind=_get_engine())
                  , autoload=True, schema=read_pgid(get_or_403(d, 'schema')))

    meta_cols = ['_message', '_user']

    field_strings = []
    for field in d.get('fields', []):
        assert ('type' in field and field['type'] == 'column'), 'Only pure column expressions are allowed in insert'

        field_strings.append(parse_expression(field))

    query = table.insert()

    if not 'method' in d:
        d['method'] = 'values'
    if d['method'] == 'values':
        if field_strings:
            raw_values = get_or_403(d, 'values')
            assert (isinstance(raw_values, list))
            values = map(lambda x: zip(field_strings, x), raw_values)
        else:
            values = get_or_403(d, 'values')

        def clear_meta(vals):
            val_dict = vals
            # make sure meta fields are not compromised
            if context['user'].is_anonymous:
                username = 'Anonymous'
            else:
                username = context['user'].name
            val_dict.update(set_meta_info('insert', username, message))
            return val_dict

        values = list(map(clear_meta, values))

        query = query.values(values)

    if 'returning' in d:
        query = query.returning(*map(Column, d['returning']))

    return query, values

def parse_select(d):
    """
        Defintion of a select query according to
        http://www.postgresql.org/docs/9.3/static/sql-select.html

        not implemented:
            [ WITH [ RECURSIVE ] with_query [, ...] ]
            [ WINDOW window_name AS ( window_definition ) [, ...] ]
            [ FOR { UPDATE | NO KEY UPDATE | SHARE | KEY SHARE } [ OF table_name [, ...] ] [ NOWAIT ] [...] ]
    """
    distinct = d.get('distinct', False)

    L = None

    if 'fields' in d and d['fields']:
        L = []
        for field in d['fields']:
            col = parse_expression(field)
            if 'as' in field:
                col.label(read_pgid(field['as']))
            L.append(col)
    from_clause = parse_from_item(get_or_403(d, 'from'))
    if not L:
        L = '*'
    query = select(columns=L, distinct=distinct, from_obj=from_clause)

    # [ WHERE condition ]
    if d.get('where', False):
        query = query.where(parse_condition(d['where']))

    if 'group_by' in d:
        query = query.group_by([parse_expression(f) for f in d['group_by']])

    if 'having' in d:
        query.having([parse_condition(f) for f in d['having']])

    if 'select' in d:
        for constraint in d['select']:
            type = get_or_403(constraint, 'type')
            subquery = parse_select(get_or_403(constraint, 'query'))
            if type.lower() == 'union':
                query.union(subquery)
            elif type.lower() == 'intersect':
                query.intersect(subquery)
            elif type.lower() == 'except':
                query.except_(subquery)
    if 'order_by' in d:
        for ob in d['order_by']:
            expr = parse_expression(ob)
            desc = ob.get('ordering', 'asc').lower() == 'desc'
            if desc:
                expr = expr.desc()
            query = query.order_by(expr)

    if 'limit' in d:
        if isinstance(d['limit'], int) or d['limit'].isdigit():
            query = query.limit(int(d['limit']))
        else:
            raise APIError('Invalid LIMIT: Expected a digit')

    if 'offset' in d:
        if isinstance(d['offset'], int) or d['offset'].isdigit():
            query = query.offset(int(d['offset']))
        else:
            raise APIError('Invalid LIMIT: Expected a digit')
    return query

def parse_from_item(d):
    """
        Defintion of a from_item according to
        http://www.postgresql.org/docs/9.3/static/sql-select.html

        return: A from_item string with checked psql qualifiers.

        Not implemented:
            with_query_name [ [ AS ] alias [ ( column_alias [, ...] ) ] ]
            [ LATERAL ] function_name ( [ argument [, ...] ] ) [ AS ] alias [ ( column_alias [, ...] | column_definition [, ...] ) ]
            [ LATERAL ] function_name ( [ argument [, ...] ] ) AS ( column_definition [, ...] )
    """
    # TODO: If 'type' is not set assume just a table name is present
    if isinstance(d, str):
        d = {'type': 'table', 'table': d}
    if isinstance(d, list):
        return [parse_from_item(f) for f in d]
    dtype = get_or_403(d, 'type')
    if dtype == 'table':
        schema_name = read_pgid(d['schema']) if 'schema' in d else None
        only = d.get('only', False)
        table_name = read_pgid(get_or_403(d, 'table'))
        item = Table(table_name, MetaData(bind=_get_engine()), schema=schema_name)
        engine = _get_engine()
        conn = engine.connect()
        exists = engine.dialect.has_table(conn, item.name, item.schema)
        conn.close()
        if not exists:
            raise APIError('Table not found: ' + str(item), status=400)
    elif dtype == 'select':
        item = parse_select(d['query'])
    elif dtype == 'join':
        left = parse_from_item(get_or_403(d, 'left'))
        right = parse_from_item(get_or_403(d, 'right'))
        is_outer = d.get('is_outer', False)
        full = d.get('is_full', False)
        on_clause = None
        if 'on' in d:
            on_clause = parse_condition(d['on'])
        item = left.join(right, onclause=on_clause, isouter=is_outer, full=full)
    else:
        raise APIError('Unknown from-item: ' + dtype)

    if 'alias' in d:
        item = item.alias(read_pgid(d['alias']))
    return item

def parse_expression(d):
    # TODO: Implement
    if isinstance(d, dict):
        dtype = get_or_403(d, 'type')
        if dtype == 'column':
            name = get_or_403(d, 'column')

            if 'table' in d:
                name = d['table'] + '.' + name
                if 'schema' in d:
                    name = d['schema'] + '.' + name
                return literal_column(name)
            return column(name)
        if dtype == 'grouping':
            return list(map(parse_expression, get_or_403(d, 'grouping')))
        if dtype == 'operator':
            return parse_operator(d)
        if dtype == 'modifier':
            return parse_modifier(d)
        if dtype == 'function':
            return parse_function(d)
        if dtype == 'star':
            return '*'
        if dtype == 'value':
            if 'value' in d:
                return read_pgvalue(get_or_403(d, 'value'))
            else:
                return None
        else:
            raise APIError('Unknown expression type: ' + dtype )
    if isinstance(d, list):
        return [parse_expression(x) for x in d]
    return d

def parse_condition(dl):
    if isinstance(dl, list):
        dl = {'type':'operator',
              'operator': 'AND',
              'operands': list(dl)}
    return parse_expression(dl)

def parse_operator(d):
    query = parse_sqla_operator(get_or_403(d, 'operator'), *list(map(parse_expression, get_or_403(d, 'operands'))))
    return query

def parse_modifier(d):
    return "%s %s" % (parse_expression(get_or_403(d,'operand')),
                      read_operator(get_or_403(d, 'operator'),
                                    get_or_403(d,'operand')))

def parse_function(d):
    fname = get_or_403(d, 'function')

    operand_struc = get_or_403(d, 'operands')
    if isinstance(operand_struc, list):
        operands = list(map(parse_expression, operand_struc))
    else:
        if isinstance(operand_struc, dict) and operand_struc.get('type',None) == 'grouping':
            operands = parse_expression(operand_struc)
        else:
            operands = [parse_expression(operand_struc)]

    if fname == '+':
        if len(operands) != 2:
            raise APIError('Wrong number of arguments for function %s. Expected 2. Got %d'%(fname, len(operands)))
        x, y = operands
        return x + y
    else:
        function = getattr(func, fname)
        return function(*operands)

def cadd(d, key, string=None):
    if not string:
        string = key.upper() + ' '
    if d.pop(key, None):
        return string
    else:
        return ''

def parse_create_table(d):
    s = 'CREATE '
    if d.pop('global', None):
        s += 'GLOBAL '
    elif d.pop('local', None):
        s += 'LOCAL '

    s += (cadd(d, 'temp')
          + cadd(d, 'unlogged')
          + 'TABLE '
          + cadd(d, 'if_not_exists', 'IF NOT EXISTS ')
          + read_pgid(get_or_403(d,'name')))

    fieldstrings = []

    for field in d.pop('fields', []):
        fs = ''
        ftype = get_or_403(field, 'type')
        if ftype == 'column':
            fs += read_pgid(get_or_403(field,'name')) + ' '
            fs += read_pgid(get_or_403(field,'data_type')) + ' '
            collate = field.pop('collate', None)
            if collate:
                fs += read_pgid(collate)
            fs += ', '.join([parse_column_constraint(cons)
                             for cons in field.pop('constraints', [])]) + ' '
        elif ftype == 'table_constraint':
            fs += parse_table_constraint(field)
        elif ftype == 'like':
            fs += 'LIKE '

def parse_column_constraint(d):
    raise NotImplementedError
    # TODO: Implement

def parse_table_constraint(d):
    raise NotImplementedError
    # TODO: Implement

def parse_scolumnd_from_columnd(schema, table, name, column_description):
    # Migrate Postgres to Python Structures
    data_type = column_description.get('data_type')
    size = column_description.get('character_maximum_length')
    if size is not None and data_type is not None:
        data_type += "(" + str(size) + ")"

    notnull = column_description.get('is_nullable', False)

    return {'column_name': name,
            'not_null': notnull,
            'data_type': data_type,
            'new_name': column_description.get('new_name'),
            'c_schema': schema,
            'c_table': table
            }

def parse_sconstd_from_constd(schema, table, name_const, constraint_description):
    defi = constraint_description.get('definition')
    return {
        'action': None,  # {ADD, DROP}
        'constraint_type': constraint_description.get('constraint_typ'),  # {FOREIGN KEY, PRIMARY KEY, UNIQUE, CHECK}
        'constraint_name': name_const,
        'constraint_parameter': constraint_description.get('definition').split('(')[1].split(')')[0],
        # Things in Brackets, e.g. name of column
        'reference_table': defi.split('REFERENCES ')[1].split('(')[2] if 'REFERENCES' in defi else None,
        'reference_column': defi.split('(')[2].split(')')[1] if 'REFERENCES' in defi else None,
        'c_schema': schema,
        'c_table': table
    }

def replace_None_with_NULL(dictonary):
    # Replacing None with null for Database
    for key, value in dictonary.items():
        if value is None:
            dictonary[key] = 'NULL'

    return dictonary

def split(string, seperator):
    if string is None:
        return None
    else:
        return str(string).split(seperator)

def replace(string, occuring_symb, replace_symb):
    if string is None:
        return None
    else:
        return str(string).replace(occuring_symb, replace_symb)

def alchemyencoder(obj):
    """JSON encoder function for SQLAlchemy special classes."""
    if isinstance(obj, datetime.date):
        return obj.isoformat()
    elif isinstance(obj, decimal.Decimal):
        return float(obj)

sql_operators = {'EQUALS': '=',
                 'GREATER': '>',
                 'LOWER': '<',
                 'NOTEQUAL': '!=',
                 'NOTGREATER': '<=',
                 'NOTLOWER': '>=',
                 '=': '=',
                 '>': '>',
                 '<': '<',
                 '!=': '!=',
                 '<>': '!=',
                 '<=': '<=',
                 '>=': '>=',
                 }

def parse_sql_operator(key: str) -> str:
    return sql_operators.get(key)

def parse_sqla_operator(key, *operands):
    if not operands:
        raise APIError('Missing arguments for \'%s\'.' % (key))
    if key in ['AND']:
        query = and_(*operands)
        return query
    elif key in ['OR']:
        query = or_(*operands)
        return query
    elif key in ['NOT']:
        x = operands[0]
        return parse_condition(x)._not()
    else:
        if len(operands) != 2:
            raise APIError('Wrong number of arguments for \'%s\'. Expected: 2 Got: %s'%(key, len(operands)))
        x, y = operands
        if key in ['EQUALS','=']:
            return x == y
        if key in ['GREATER', '>']:
            return x > y
        if key in ['LOWER', '<']:
            return x < y
        if key in ['NOTEQUAL', '<>', '!=']:
            return x != y
        if key in ['NOTGREATER', '<=']:
            return x <= y
        if key in ['NOTLOWER', '>=']:
            return x >= y
        if key in ['as']:
            return x.label(y)

    raise APIError("Operator %s not supported"%key)<commit_msg>Make operator handling case insensitive<commit_after>###########
# Parsers #
###########
import decimal
import re
from datetime import datetime
from sqlalchemy import Table, MetaData, Column, select, column, func, literal_column, and_, or_
from api.error import APIError, APIKeyError
from api.connection import _get_engine
import geoalchemy2  # Although this import seems unused is has to be here

pgsql_qualifier = re.compile(r"^[\w\d_\.]+$")

def get_or_403(dictionary, key):
    try:
        return dictionary[key]
    except KeyError:
        raise APIKeyError(dictionary, key)

def is_pg_qual(x):
    return pgsql_qualifier.search(x)

def quote(x):
    if not x.startswith('"') and '(' not in x:
        return '"' + x + '"'
    else:
        return x

def read_pgvalue(x):
    # TODO: Implement check for valid values
    if isinstance(x, str):
        return "'" + x + "'"
    if x is None:
        return 'null'
    return x

def read_operator(x, right):
    # TODO: Implement check for valid operators
    if isinstance(right, dict) and get_or_403(right, 'type') == 'value' and ('value' not in right or get_or_403(right, 'value') is None and x == '='):
        return 'is'
    return x

class ValidationError(Exception):
    def __init__(self, message, value):
        self.message = message
        self.value = value

def read_bool(s):
    if isinstance(s, bool):
        return s
    if s.lower() in ["true", "false"]:
        return s.lower() == "true"
    elif s.lower() in ["yes", "no"]:
        return s.lower() == "true"
    else:
        raise APIError("Invalid value in binary field", s)

def read_pgid(s):
    if is_pg_qual(s):
        return s
    raise APIError("Invalid identifier: '%s'"%s)

def set_meta_info(method, user, message=None):
    val_dict = {}
    val_dict['_user'] = user  # TODO: Add user handling
    val_dict['_message'] = message
    return val_dict

def parse_insert(d, context, message=None):
    table = Table(read_pgid(get_or_403(d, 'table')), MetaData(bind=_get_engine())
                  , autoload=True, schema=read_pgid(get_or_403(d, 'schema')))

    meta_cols = ['_message', '_user']

    field_strings = []
    for field in d.get('fields', []):
        assert ('type' in field and field['type'] == 'column'), 'Only pure column expressions are allowed in insert'

        field_strings.append(parse_expression(field))

    query = table.insert()

    if not 'method' in d:
        d['method'] = 'values'
    if d['method'] == 'values':
        if field_strings:
            raw_values = get_or_403(d, 'values')
            assert (isinstance(raw_values, list))
            values = map(lambda x: zip(field_strings, x), raw_values)
        else:
            values = get_or_403(d, 'values')

        def clear_meta(vals):
            val_dict = vals
            # make sure meta fields are not compromised
            if context['user'].is_anonymous:
                username = 'Anonymous'
            else:
                username = context['user'].name
            val_dict.update(set_meta_info('insert', username, message))
            return val_dict

        values = list(map(clear_meta, values))

        query = query.values(values)

    if 'returning' in d:
        query = query.returning(*map(Column, d['returning']))

    return query, values

def parse_select(d):
    """
        Defintion of a select query according to
        http://www.postgresql.org/docs/9.3/static/sql-select.html

        not implemented:
            [ WITH [ RECURSIVE ] with_query [, ...] ]
            [ WINDOW window_name AS ( window_definition ) [, ...] ]
            [ FOR { UPDATE | NO KEY UPDATE | SHARE | KEY SHARE } [ OF table_name [, ...] ] [ NOWAIT ] [...] ]
    """
    distinct = d.get('distinct', False)

    L = None

    if 'fields' in d and d['fields']:
        L = []
        for field in d['fields']:
            col = parse_expression(field)
            if 'as' in field:
                col.label(read_pgid(field['as']))
            L.append(col)
    from_clause = parse_from_item(get_or_403(d, 'from'))
    if not L:
        L = '*'
    query = select(columns=L, distinct=distinct, from_obj=from_clause)

    # [ WHERE condition ]
    if d.get('where', False):
        query = query.where(parse_condition(d['where']))

    if 'group_by' in d:
        query = query.group_by([parse_expression(f) for f in d['group_by']])

    if 'having' in d:
        query.having([parse_condition(f) for f in d['having']])

    if 'select' in d:
        for constraint in d['select']:
            type = get_or_403(constraint, 'type')
            subquery = parse_select(get_or_403(constraint, 'query'))
            if type.lower() == 'union':
                query.union(subquery)
            elif type.lower() == 'intersect':
                query.intersect(subquery)
            elif type.lower() == 'except':
                query.except_(subquery)
    if 'order_by' in d:
        for ob in d['order_by']:
            expr = parse_expression(ob)
            desc = ob.get('ordering', 'asc').lower() == 'desc'
            if desc:
                expr = expr.desc()
            query = query.order_by(expr)

    if 'limit' in d:
        if isinstance(d['limit'], int) or d['limit'].isdigit():
            query = query.limit(int(d['limit']))
        else:
            raise APIError('Invalid LIMIT: Expected a digit')

    if 'offset' in d:
        if isinstance(d['offset'], int) or d['offset'].isdigit():
            query = query.offset(int(d['offset']))
        else:
            raise APIError('Invalid LIMIT: Expected a digit')
    return query

def parse_from_item(d):
    """
        Defintion of a from_item according to
        http://www.postgresql.org/docs/9.3/static/sql-select.html

        return: A from_item string with checked psql qualifiers.

        Not implemented:
            with_query_name [ [ AS ] alias [ ( column_alias [, ...] ) ] ]
            [ LATERAL ] function_name ( [ argument [, ...] ] ) [ AS ] alias [ ( column_alias [, ...] | column_definition [, ...] ) ]
            [ LATERAL ] function_name ( [ argument [, ...] ] ) AS ( column_definition [, ...] )
    """
    # TODO: If 'type' is not set assume just a table name is present
    if isinstance(d, str):
        d = {'type': 'table', 'table': d}
    if isinstance(d, list):
        return [parse_from_item(f) for f in d]
    dtype = get_or_403(d, 'type')
    if dtype == 'table':
        schema_name = read_pgid(d['schema']) if 'schema' in d else None
        only = d.get('only', False)
        table_name = read_pgid(get_or_403(d, 'table'))
        item = Table(table_name, MetaData(bind=_get_engine()), schema=schema_name)
        engine = _get_engine()
        conn = engine.connect()
        exists = engine.dialect.has_table(conn, item.name, item.schema)
        conn.close()
        if not exists:
            raise APIError('Table not found: ' + str(item), status=400)
    elif dtype == 'select':
        item = parse_select(d['query'])
    elif dtype == 'join':
        left = parse_from_item(get_or_403(d, 'left'))
        right = parse_from_item(get_or_403(d, 'right'))
        is_outer = d.get('is_outer', False)
        full = d.get('is_full', False)
        on_clause = None
        if 'on' in d:
            on_clause = parse_condition(d['on'])
        item = left.join(right, onclause=on_clause, isouter=is_outer, full=full)
    else:
        raise APIError('Unknown from-item: ' + dtype)

    if 'alias' in d:
        item = item.alias(read_pgid(d['alias']))
    return item

def parse_expression(d):
    # TODO: Implement
    if isinstance(d, dict):
        dtype = get_or_403(d, 'type')
        if dtype == 'column':
            name = get_or_403(d, 'column')

            if 'table' in d:
                name = d['table'] + '.' + name
                if 'schema' in d:
                    name = d['schema'] + '.' + name
                return literal_column(name)
            return column(name)
        if dtype == 'grouping':
            return list(map(parse_expression, get_or_403(d, 'grouping')))
        if dtype == 'operator':
            return parse_operator(d)
        if dtype == 'modifier':
            return parse_modifier(d)
        if dtype == 'function':
            return parse_function(d)
        if dtype == 'star':
            return '*'
        if dtype == 'value':
            if 'value' in d:
                return read_pgvalue(get_or_403(d, 'value'))
            else:
                return None
        else:
            raise APIError('Unknown expression type: ' + dtype )
    if isinstance(d, list):
        return [parse_expression(x) for x in d]
    return d

def parse_condition(dl):
    if isinstance(dl, list):
        dl = {'type':'operator',
              'operator': 'AND',
              'operands': list(dl)}
    return parse_expression(dl)

def parse_operator(d):
    query = parse_sqla_operator(get_or_403(d, 'operator'), *list(map(parse_expression, get_or_403(d, 'operands'))))
    return query

def parse_modifier(d):
    return "%s %s" % (parse_expression(get_or_403(d,'operand')),
                      read_operator(get_or_403(d, 'operator'),
                                    get_or_403(d,'operand')))

def parse_function(d):
    fname = get_or_403(d, 'function')

    operand_struc = get_or_403(d, 'operands')
    if isinstance(operand_struc, list):
        operands = list(map(parse_expression, operand_struc))
    else:
        if isinstance(operand_struc, dict) and operand_struc.get('type',None) == 'grouping':
            operands = parse_expression(operand_struc)
        else:
            operands = [parse_expression(operand_struc)]

    if fname == '+':
        if len(operands) != 2:
            raise APIError('Wrong number of arguments for function %s. Expected 2. Got %d'%(fname, len(operands)))
        x, y = operands
        return x + y
    else:
        function = getattr(func, fname)
        return function(*operands)

def cadd(d, key, string=None):
    if not string:
        string = key.upper() + ' '
    if d.pop(key, None):
        return string
    else:
        return ''

def parse_create_table(d):
    s = 'CREATE '
    if d.pop('global', None):
        s += 'GLOBAL '
    elif d.pop('local', None):
        s += 'LOCAL '

    s += (cadd(d, 'temp')
          + cadd(d, 'unlogged')
          + 'TABLE '
          + cadd(d, 'if_not_exists', 'IF NOT EXISTS ')
          + read_pgid(get_or_403(d,'name')))

    fieldstrings = []

    for field in d.pop('fields', []):
        fs = ''
        ftype = get_or_403(field, 'type')
        if ftype == 'column':
            fs += read_pgid(get_or_403(field,'name')) + ' '
            fs += read_pgid(get_or_403(field,'data_type')) + ' '
            collate = field.pop('collate', None)
            if collate:
                fs += read_pgid(collate)
            fs += ', '.join([parse_column_constraint(cons)
                             for cons in field.pop('constraints', [])]) + ' '
        elif ftype == 'table_constraint':
            fs += parse_table_constraint(field)
        elif ftype == 'like':
            fs += 'LIKE '

def parse_column_constraint(d):
    raise NotImplementedError
    # TODO: Implement

def parse_table_constraint(d):
    raise NotImplementedError
    # TODO: Implement

def parse_scolumnd_from_columnd(schema, table, name, column_description):
    # Migrate Postgres to Python Structures
    data_type = column_description.get('data_type')
    size = column_description.get('character_maximum_length')
    if size is not None and data_type is not None:
        data_type += "(" + str(size) + ")"

    notnull = column_description.get('is_nullable', False)

    return {'column_name': name,
            'not_null': notnull,
            'data_type': data_type,
            'new_name': column_description.get('new_name'),
            'c_schema': schema,
            'c_table': table
            }

def parse_sconstd_from_constd(schema, table, name_const, constraint_description):
    defi = constraint_description.get('definition')
    return {
        'action': None,  # {ADD, DROP}
        'constraint_type': constraint_description.get('constraint_typ'),  # {FOREIGN KEY, PRIMARY KEY, UNIQUE, CHECK}
        'constraint_name': name_const,
        'constraint_parameter': constraint_description.get('definition').split('(')[1].split(')')[0],
        # Things in Brackets, e.g. name of column
        'reference_table': defi.split('REFERENCES ')[1].split('(')[2] if 'REFERENCES' in defi else None,
        'reference_column': defi.split('(')[2].split(')')[1] if 'REFERENCES' in defi else None,
        'c_schema': schema,
        'c_table': table
    }

def replace_None_with_NULL(dictonary):
    # Replacing None with null for Database
    for key, value in dictonary.items():
        if value is None:
            dictonary[key] = 'NULL'

    return dictonary

def split(string, seperator):
    if string is None:
        return None
    else:
        return str(string).split(seperator)

def replace(string, occuring_symb, replace_symb):
    if string is None:
        return None
    else:
        return str(string).replace(occuring_symb, replace_symb)

def alchemyencoder(obj):
    """JSON encoder function for SQLAlchemy special classes."""
    if isinstance(obj, datetime.date):
        return obj.isoformat()
    elif isinstance(obj, decimal.Decimal):
        return float(obj)

sql_operators = {'EQUALS': '=',
                 'GREATER': '>',
                 'LOWER': '<',
                 'NOTEQUAL': '!=',
                 'NOTGREATER': '<=',
                 'NOTLOWER': '>=',
                 '=': '=',
                 '>': '>',
                 '<': '<',
                 '!=': '!=',
                 '<>': '!=',
                 '<=': '<=',
                 '>=': '>=',
                 }

def parse_sql_operator(key: str) -> str:
    return sql_operators.get(key)

def parse_sqla_operator(raw_key, *operands):
    key = raw_key.lower().strip()
    if not operands:
        raise APIError('Missing arguments for \'%s\'.' % (key))
    if key in ['and']:
        query = and_(*operands)
        return query
    elif key in ['or']:
        query = or_(*operands)
        return query
    elif key in ['not']:
        x = operands[0]
        return parse_condition(x)._not()
    else:
        if len(operands) != 2:
            raise APIError('Wrong number of arguments for \'%s\'. Expected: 2 Got: %s'%(key, len(operands)))
        x, y = operands
        if key in ['equals','=']:
            return x == y
        if key in ['greater', '>']:
            return x > y
        if key in ['lower', '<']:
            return x < y
        if key in ['notequal', '<>', '!=']:
            return x != y
        if key in ['notgreater', '<=']:
            return x <= y
        if key in ['notlower', '>=']:
            return x >= y
        if key in ['as']:
            return x.label(y)

    raise APIError("Operator %s not supported"%key)<|endoftext|>


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

# Hive Netius System
# Copyright (c) 2008-2017 Hive Solutions Lda.
#
# This file is part of Hive Netius System.
#
# Hive Netius System is free software: you can redistribute it and/or modify
# it under the terms of the Apache License as published by the Apache
# Foundation, either version 2.0 of the License, or (at your option) any
# later version.
#
# Hive Netius System is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# Apache License for more details.
#
# You should have received a copy of the Apache License along with
# Hive Netius System. If not, see <http://www.apache.org/licenses/>.

__author__ = "João Magalhães <joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2017 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "Apache License, Version 2.0"
""" The license for the module """

import math
import base64

import netius

from . import asn
from . import util
from . import calc

PRIVATE_TOKEN = "RSA PRIVATE KEY"
PUBLIC_TOKEN = "PUBLIC KEY"

def open_pem_key(path, token = PRIVATE_TOKEN):
    is_file = not type(path) in netius.legacy.STRINGS
    if is_file: file = path
    else: file = open(path, "rb")
    try:
        data = file.read()
    finally:
        if not is_file: file.close()
    return open_pem_data(data, token = token)

def open_pem_data(data, token = PRIVATE_TOKEN):
    begin, end = pem_limiters(token)

    begin_index = data.find(begin)
    end_index = data.find(end)

    if begin_index == -1: raise netius.ParserError("Invalid key format")
    if end_index == -1: raise netius.ParserError("Invalid key format")

    begin_index += len(begin)

    data = data[begin_index:end_index]
    data = data.strip()
    return base64.b64decode(data)

def write_pem_key(
    path,
    data,
    token = PRIVATE_TOKEN,
    width = 64
):
    begin, end = pem_limiters(token)

    data = base64.b64encode(data)

    chunks = [chunk for chunk in util.chunks(data, width)]
    data = b"\n".join(chunks)

    is_file = not type(path) in netius.legacy.STRINGS
    file = path if is_file else open(path, "wb")
    try:
        file.write(begin)
        file.write(b"\n")
        file.write(data)
        file.write(b"\n")
        file.write(end)
        file.write(b"\n")
    finally:
        if not is_file: file.close()

def open_private_key(path):
    data = open_pem_key(
        path,
        token = PRIVATE_TOKEN
    )
    return open_private_key_data(data)

def open_private_key_b64(data_b64):
    data = base64.b64decode(data_b64)
    return open_private_key_data(data)

def open_private_key_data(data):
    asn1 = asn.asn1_parse(asn.ASN1_RSA_PRIVATE_KEY, data)[0]
    private_key = dict(
        version = asn1[0],
        modulus = asn1[1],
        public_exponent = asn1[2],
        private_exponent = asn1[3],
        prime_1 = asn1[4],
        prime_2 = asn1[5],
        exponent_1 = asn1[6],
        exponent_2 = asn1[7],
        coefficient = asn1[8],
        bits = rsa_bits(asn1[1])
    )
    return private_key

def open_public_key(path):
    data = open_pem_key(
        path,
        token = PUBLIC_TOKEN
    )
    return open_public_key_data(data)

def open_public_key_b64(data_b64):
    data = base64.b64decode(data_b64)
    return open_public_key_data(data)

def open_public_key_data(data):
    asn1 = asn.asn1_parse(asn.ASN1_OBJECT, data)[0]
    asn1 = asn.asn1_parse(asn.ASN1_RSA_PUBLIC_KEY, asn1[1][1:])[0]
    public_key = dict(
        modulus = asn1[0],
        public_exponent = asn1[1],
        bits = rsa_bits(asn1[0])
    )
    return public_key

def write_private_key(path, private_key):
    data = asn_private_key(private_key)
    write_pem_key(
        path,
        data,
        token = PRIVATE_TOKEN
    )

def write_public_key(path, public_key):
    data = asn_public_key(public_key)
    write_pem_key(
        path,
        data,
        token = PUBLIC_TOKEN
    )

def asn_private_key(private_key):
    return asn.asn1_gen(
        (asn.SEQUENCE, [
            (asn.INTEGER, private_key["version"]),
            (asn.INTEGER, private_key["modulus"]),
            (asn.INTEGER, private_key["public_exponent"]),
            (asn.INTEGER, private_key["private_exponent"]),
            (asn.INTEGER, private_key["prime_1"]),
            (asn.INTEGER, private_key["prime_2"]),
            (asn.INTEGER, private_key["exponent_1"]),
            (asn.INTEGER, private_key["exponent_2"]),
            (asn.INTEGER, private_key["coefficient"])
        ])
    )

def asn_public_key(public_key):
    data = b"\x00" + asn.asn1_gen(
        (asn.SEQUENCE, [
            (asn.INTEGER, public_key["modulus"]),
            (asn.INTEGER, public_key["public_exponent"])
        ])
    )
    return asn.asn1_gen(
        (asn.SEQUENCE, [
            (asn.SEQUENCE, [
                (asn.OBJECT_IDENTIFIER, asn.RSAID_PKCS1),
                (asn.NULL, None)
            ]),
            (asn.BIT_STRING, data)
        ])
    )

def pem_to_der(in_path, out_path, token = PRIVATE_TOKEN):
    data = open_pem_key(in_path, token = token)
    file = open(out_path, "wb")
    try: file.write(data)
    finally: file.close()

def pem_limiters(token):
    begin = netius.legacy.bytes("-----BEGIN " + token + "-----")
    end = netius.legacy.bytes("-----END " + token + "-----")
    return (begin, end)

def private_to_public(private_key):
    public_key = dict(
        modulus = private_key["modulus"],
        public_exponent = private_key["public_exponent"]
    )
    return public_key

def assert_private(private_key, number_bits = None):
    prime_1 = private_key["prime_1"]
    prime_2 = private_key["prime_2"]
    private_exponent = private_key["private_exponent"]

    modulus = prime_1 * prime_2
    exponent_1 = private_exponent % (prime_1 - 1)
    exponent_2 = private_exponent % (prime_2 - 1)
    coefficient = calc.modinv(prime_2, prime_1)

    util.verify(modulus == private_key["modulus"])
    util.verify(exponent_1 == private_key["exponent_1"])
    util.verify(exponent_2 == private_key["exponent_2"])
    util.verify(coefficient == private_key["coefficient"])

    if number_bits:
        util.verify(number_bits // 2 == rsa_bits(private_key["prime_1"]))
        util.verify(number_bits // 2 == rsa_bits(private_key["prime_2"]))
        util.verify(number_bits == private_key["bits"])

    message = b"Hello World"
    signature = rsa_sign(b"Hello World", private_key)
    result = rsa_verify(signature, private_key)
    result = result.lstrip(b"\0")

    util.verify(result == message)

def rsa_private(number_bits):
    """
    Generates a new "random" private with the requested number
    of bits as the base for exponents and modulus.

    This method is extremely time consuming in terms of processor
    and should be used carefully to avoid any problem.

    :type number_bits: int
    :param number_bits: The number of bits that are going to be
    used for the generation of the private key.
    :rtype: Dictionary
    :return: The generated private key structure, may then be used
    for processing or written to a file.
    """

    while True:
        prime_1, prime_2 = rsa_primes(number_bits // 2)
        public_exponent, private_exponent = rsa_exponents(prime_1, prime_2, number_bits // 2)
        if private_exponent > 0: break

    modulus = prime_1 * prime_2
    exponent_1 = private_exponent % (prime_1 - 1)
    exponent_2 = private_exponent % (prime_2 - 1)
    coefficient = calc.modinv(prime_2, prime_1)
    bits = rsa_bits(modulus)

    private_key = dict(
        version = 0,
        modulus = modulus,
        public_exponent = public_exponent,
        private_exponent = private_exponent,
        prime_1 = prime_1,
        prime_2 = prime_2,
        exponent_1 = exponent_1,
        exponent_2 = exponent_2,
        coefficient = coefficient,
        bits = bits
    )

    return private_key

def rsa_primes(number_bits):
    """
    Generates two different prime numbers (p and q values)
    and returns them inside a tuple structure.

    The generation is made according to the number of bits
    defined and using a trial and error strategy (expensive).

    :type number_bits: int
    :param number_bits: The number of bits to be used in
    prime generation, this affects security.
    :rtype: Tuple
    :return: A tuple containing the two different prime
    numbers to be returned.
    """

    # calculates the total number of bits for the key as
    # the double of the requested for the primes generation
    total_bits = number_bits * 2

    # constructs the clojure based function that is going to
    # be used as the validator of the primes combination, this
    # used for a trial an error based approach for the generation
    # of the primes to be used in the private key
    def rsa_acceptable(prime_1, prime_2):
        if prime_1 == prime_2: return False
        modulus_bits = rsa_bits(prime_1 * prime_2)
        return modulus_bits == total_bits

    # generates the "first" version of both prime values
    # that is going to serve as the first iteration of
    # each of the values, and start the is odd variable
    prime_1 = calc.prime(number_bits)
    prime_2 = calc.prime(number_bits)
    is_odd = True

    # iterates continuously trying to find a combination
    # of prime numbers that is acceptable and valid
    while True:
        if rsa_acceptable(prime_1, prime_2): break
        if is_odd: prime_1 = calc.prime(number_bits)
        else: prime_2 = calc.prime(number_bits)

    # returns a tuple containing both of the generated
    # primes and returns it to the caller method
    return (prime_1, prime_2)

def rsa_exponents(prime_1, prime_2, number_bits, basic = True):
    """
    Generates both the public and the private exponents for
    the rsa cryptography system taking as base the provided
    prime numbers and the amount of bits for the values.

    :type prime_1: int
    :param prime_1: The first prime number use for rsa.
    :type prime_2: int
    :param prime_2: The second prime number use for rsa.
    :type number_bits: int
    :param number_bits: The number of bits that are going to be
    used for the generation of the values.
    :type basic: bool
    :param basic: If the basic approach to the generation of the
    public exponent should be taken into account.
    :rtype: Tuple
    :return: The tuple containing the generated public and
    private keys (properly tested).
    """

    # calculates the modulus and the phi value for the
    # modulus, as the y are going to be used for calculus
    modulus = prime_1 * prime_2
    phi_modulus = (prime_1 - 1) * (prime_2 - 1)

    # starts by setting the is first flag so that the first iteration
    # of the public exponent generation cycle is taking into account
    # the possible basic flag value for the public exponent
    is_first = True

    # iterates continuously to find a valid public exponent, one
    # that satisfies the relative prime
    while True:
        # make sure e has enough bits so we ensure "wrapping" through
        # modulus (n value) note that if the this is the first attempt
        # to create a public exponent and the basic mode is active the
        # number chosen is the "magic" number (compatibility)
        public_exponent = calc.prime(max(8, number_bits // 2))
        if is_first and basic: public_exponent = 65537; is_first = False

        # checks if the exponent and the modulus are relative primes
        # and also checks if the exponent and the phi modulus are relative
        # primes, for that situation a valid public exponent has been fond
        # and the cycle may be broken
        is_relative = calc.relatively_prime(public_exponent, modulus)
        is_relative_phi = calc.relatively_prime(public_exponent, phi_modulus)
        if is_relative and is_relative_phi: break

    # retrieves the result of the extended euclid greatest common divisor,
    # this value is going to be used as the basis for the calculus of the
    # private exponent for the current operation
    d, l, _e = calc.egcd(public_exponent, phi_modulus)
    private_exponent = l

    # in case the greatest common divisor between both is not one, the values
    # are not relative primes and an exception must be raised
    if not d == 1: raise netius.GeneratorError(
        "The public exponent '%d' and the phi modulus '%d' are not relative primes" %
        (public_exponent, phi_modulus)
    )

    # calculates the inverse modulus for both exponent and in case it's not one
    # an exception is raised about the problem
    inverse_modulus = (public_exponent * private_exponent) % phi_modulus
    if not inverse_modulus == 1: netius.GeneratorError(
        "The public exponent '%d' and private exponent '%d' are not multiplicative inverse modulus of phi modulus '%d'" %
        (public_exponent, private_exponent, phi_modulus)
    )

    # creates the tuple that contains both the public and the private
    # exponent values that may be used for rsa based cryptography
    return (public_exponent, private_exponent)

def rsa_bits(modulus):
    bits = math.log(modulus, 2)
    return calc.ceil_integer(bits)

def rsa_sign(message, private_key):
    message = netius.legacy.bytes(message)
    modulus = private_key["modulus"]
    private_exponent = private_key["private_exponent"]
    return rsa_crypt_s(message, private_exponent, modulus)

def rsa_verify(signature, public_key):
    signature = netius.legacy.bytes(signature)
    modulus = public_key["modulus"]
    public_exponent = public_key["public_exponent"]
    return rsa_crypt_s(signature, public_exponent, modulus)

def rsa_crypt_s(message, exponent, modulus):
    modulus_l = calc.ceil_integer(math.log(modulus, 256))

    message_i = util.bytes_to_integer(message)
    message_crypt = rsa_crypt(message_i, exponent, modulus)
    message_crypt_s = util.integer_to_bytes(message_crypt, modulus_l)

    return message_crypt_s

def rsa_crypt(number, exponent, modulus):
    if not type(number) in netius.legacy.INTEGERS:
        raise TypeError("you must pass a long or an int")

    if number > 0 and math.floor(math.log(number, 2)) > math.floor(math.log(modulus, 2)):
        raise OverflowError("the message is too long")

    return pow(number, exponent, modulus)
<commit_msg>small fixes<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# Hive Netius System
# Copyright (c) 2008-2017 Hive Solutions Lda.
#
# This file is part of Hive Netius System.
#
# Hive Netius System is free software: you can redistribute it and/or modify
# it under the terms of the Apache License as published by the Apache
# Foundation, either version 2.0 of the License, or (at your option) any
# later version.
#
# Hive Netius System is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# Apache License for more details.
#
# You should have received a copy of the Apache License along with
# Hive Netius System. If not, see <http://www.apache.org/licenses/>.

__author__ = "João Magalhães <joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2017 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "Apache License, Version 2.0"
""" The license for the module """

import math
import base64

import netius

from . import asn
from . import util
from . import calc

PRIVATE_TOKEN = "RSA PRIVATE KEY"
PUBLIC_TOKEN = "PUBLIC KEY"

def open_pem_key(path, token = PRIVATE_TOKEN):
    is_file = not isinstance(path, netius.legacy.STRINGS)
    if is_file: file = path
    else: file = open(path, "rb")
    try:
        data = file.read()
    finally:
        if not is_file: file.close()
    return open_pem_data(data, token = token)

def open_pem_data(data, token = PRIVATE_TOKEN):
    begin, end = pem_limiters(token)

    begin_index = data.find(begin)
    end_index = data.find(end)

    if begin_index == -1: raise netius.ParserError("Invalid key format")
    if end_index == -1: raise netius.ParserError("Invalid key format")

    begin_index += len(begin)

    data = data[begin_index:end_index]
    data = data.strip()
    return base64.b64decode(data)

def write_pem_key(
    path,
    data,
    token = PRIVATE_TOKEN,
    width = 64
):
    begin, end = pem_limiters(token)

    data = base64.b64encode(data)

    chunks = [chunk for chunk in util.chunks(data, width)]
    data = b"\n".join(chunks)

    is_file = not isinstance(path, netius.legacy.STRINGS)
    file = path if is_file else open(path, "wb")
    try:
        file.write(begin)
        file.write(b"\n")
        file.write(data)
        file.write(b"\n")
        file.write(end)
        file.write(b"\n")
    finally:
        if not is_file: file.close()

def open_private_key(path):
    data = open_pem_key(
        path,
        token = PRIVATE_TOKEN
    )
    return open_private_key_data(data)

def open_private_key_b64(data_b64):
    data = base64.b64decode(data_b64)
    return open_private_key_data(data)

def open_private_key_data(data):
    asn1 = asn.asn1_parse(asn.ASN1_RSA_PRIVATE_KEY, data)[0]
    private_key = dict(
        version = asn1[0],
        modulus = asn1[1],
        public_exponent = asn1[2],
        private_exponent = asn1[3],
        prime_1 = asn1[4],
        prime_2 = asn1[5],
        exponent_1 = asn1[6],
        exponent_2 = asn1[7],
        coefficient = asn1[8],
        bits = rsa_bits(asn1[1])
    )
    return private_key

def open_public_key(path):
    data = open_pem_key(
        path,
        token = PUBLIC_TOKEN
    )
    return open_public_key_data(data)

def open_public_key_b64(data_b64):
    data = base64.b64decode(data_b64)
    return open_public_key_data(data)

def open_public_key_data(data):
    asn1 = asn.asn1_parse(asn.ASN1_OBJECT, data)[0]
    asn1 = asn.asn1_parse(asn.ASN1_RSA_PUBLIC_KEY, asn1[1][1:])[0]
    public_key = dict(
        modulus = asn1[0],
        public_exponent = asn1[1],
        bits = rsa_bits(asn1[0])
    )
    return public_key

def write_private_key(path, private_key):
    data = asn_private_key(private_key)
    write_pem_key(
        path,
        data,
        token = PRIVATE_TOKEN
    )

def write_public_key(path, public_key):
    data = asn_public_key(public_key)
    write_pem_key(
        path,
        data,
        token = PUBLIC_TOKEN
    )

def asn_private_key(private_key):
    return asn.asn1_gen(
        (asn.SEQUENCE, [
            (asn.INTEGER, private_key["version"]),
            (asn.INTEGER, private_key["modulus"]),
            (asn.INTEGER, private_key["public_exponent"]),
            (asn.INTEGER, private_key["private_exponent"]),
            (asn.INTEGER, private_key["prime_1"]),
            (asn.INTEGER, private_key["prime_2"]),
            (asn.INTEGER, private_key["exponent_1"]),
            (asn.INTEGER, private_key["exponent_2"]),
            (asn.INTEGER, private_key["coefficient"])
        ])
    )

def asn_public_key(public_key):
    data = b"\x00" + asn.asn1_gen(
        (asn.SEQUENCE, [
            (asn.INTEGER, public_key["modulus"]),
            (asn.INTEGER, public_key["public_exponent"])
        ])
    )
    return asn.asn1_gen(
        (asn.SEQUENCE, [
            (asn.SEQUENCE, [
                (asn.OBJECT_IDENTIFIER, asn.RSAID_PKCS1),
                (asn.NULL, None)
            ]),
            (asn.BIT_STRING, data)
        ])
    )

def pem_to_der(in_path, out_path, token = PRIVATE_TOKEN):
    data = open_pem_key(in_path, token = token)
    file = open(out_path, "wb")
    try: file.write(data)
    finally: file.close()

def pem_limiters(token):
    begin = netius.legacy.bytes("-----BEGIN " + token + "-----")
    end = netius.legacy.bytes("-----END " + token + "-----")
    return (begin, end)

def private_to_public(private_key):
    public_key = dict(
        modulus = private_key["modulus"],
        public_exponent = private_key["public_exponent"]
    )
    return public_key

def assert_private(private_key, number_bits = None):
    prime_1 = private_key["prime_1"]
    prime_2 = private_key["prime_2"]
    private_exponent = private_key["private_exponent"]

    modulus = prime_1 * prime_2
    exponent_1 = private_exponent % (prime_1 - 1)
    exponent_2 = private_exponent % (prime_2 - 1)
    coefficient = calc.modinv(prime_2, prime_1)

    util.verify(modulus == private_key["modulus"])
    util.verify(exponent_1 == private_key["exponent_1"])
    util.verify(exponent_2 == private_key["exponent_2"])
    util.verify(coefficient == private_key["coefficient"])

    if number_bits:
        util.verify(number_bits // 2 == rsa_bits(private_key["prime_1"]))
        util.verify(number_bits // 2 == rsa_bits(private_key["prime_2"]))
        util.verify(number_bits == private_key["bits"])

    message = b"Hello World"
    signature = rsa_sign(b"Hello World", private_key)
    result = rsa_verify(signature, private_key)
    result = result.lstrip(b"\0")

    util.verify(result == message)

def rsa_private(number_bits):
    """
    Generates a new "random" private with the requested number
    of bits as the base for exponents and modulus.

    This method is extremely time consuming in terms of processor
    and should be used carefully to avoid any problem.

    :type number_bits: int
    :param number_bits: The number of bits that are going to be
    used for the generation of the private key.
    :rtype: Dictionary
    :return: The generated private key structure, may then be used
    for processing or written to a file.
    """

    while True:
        prime_1, prime_2 = rsa_primes(number_bits // 2)
        public_exponent, private_exponent = rsa_exponents(prime_1, prime_2, number_bits // 2)
        if private_exponent > 0: break

    modulus = prime_1 * prime_2
    exponent_1 = private_exponent % (prime_1 - 1)
    exponent_2 = private_exponent % (prime_2 - 1)
    coefficient = calc.modinv(prime_2, prime_1)
    bits = rsa_bits(modulus)

    private_key = dict(
        version = 0,
        modulus = modulus,
        public_exponent = public_exponent,
        private_exponent = private_exponent,
        prime_1 = prime_1,
        prime_2 = prime_2,
        exponent_1 = exponent_1,
        exponent_2 = exponent_2,
        coefficient = coefficient,
        bits = bits
    )

    return private_key

def rsa_primes(number_bits):
    """
    Generates two different prime numbers (p and q values)
    and returns them inside a tuple structure.

    The generation is made according to the number of bits
    defined and using a trial and error strategy (expensive).

    :type number_bits: int
    :param number_bits: The number of bits to be used in
    prime generation, this affects security.
    :rtype: Tuple
    :return: A tuple containing the two different prime
    numbers to be returned.
    """

    # calculates the total number of bits for the key as
    # the double of the requested for the primes generation
    total_bits = number_bits * 2

    # constructs the clojure based function that is going to
    # be used as the validator of the primes combination, this
    # used for a trial an error based approach for the generation
    # of the primes to be used in the private key
    def rsa_acceptable(prime_1, prime_2):
        if prime_1 == prime_2: return False
        modulus_bits = rsa_bits(prime_1 * prime_2)
        return modulus_bits == total_bits

    # generates the "first" version of both prime values
    # that is going to serve as the first iteration of
    # each of the values, and start the is odd variable
    prime_1 = calc.prime(number_bits)
    prime_2 = calc.prime(number_bits)
    is_odd = True

    # iterates continuously trying to find a combination
    # of prime numbers that is acceptable and valid
    while True:
        if rsa_acceptable(prime_1, prime_2): break
        if is_odd: prime_1 = calc.prime(number_bits)
        else: prime_2 = calc.prime(number_bits)

    # returns a tuple containing both of the generated
    # primes and returns it to the caller method
    return (prime_1, prime_2)

def rsa_exponents(prime_1, prime_2, number_bits, basic = True):
    """
    Generates both the public and the private exponents for
    the rsa cryptography system taking as base the provided
    prime numbers and the amount of bits for the values.

    :type prime_1: int
    :param prime_1: The first prime number use for rsa.
    :type prime_2: int
    :param prime_2: The second prime number use for rsa.
    :type number_bits: int
    :param number_bits: The number of bits that are going to be
    used for the generation of the values.
    :type basic: bool
    :param basic: If the basic approach to the generation of the
    public exponent should be taken into account.
    :rtype: Tuple
    :return: The tuple containing the generated public and
    private keys (properly tested).
    """

    # calculates the modulus and the phi value for the
    # modulus, as the y are going to be used for calculus
    modulus = prime_1 * prime_2
    phi_modulus = (prime_1 - 1) * (prime_2 - 1)

    # starts by setting the is first flag so that the first iteration
    # of the public exponent generation cycle is taking into account
    # the possible basic flag value for the public exponent
    is_first = True

    # iterates continuously to find a valid public exponent, one
    # that satisfies the relative prime
    while True:
        # make sure e has enough bits so we ensure "wrapping" through
        # modulus (n value) note that if the this is the first attempt
        # to create a public exponent and the basic mode is active the
        # number chosen is the "magic" number (compatibility)
        public_exponent = calc.prime(max(8, number_bits // 2))
        if is_first and basic: public_exponent = 65537; is_first = False

        # checks if the exponent and the modulus are relative primes
        # and also checks if the exponent and the phi modulus are relative
        # primes, for that situation a valid public exponent has been fond
        # and the cycle may be broken
        is_relative = calc.relatively_prime(public_exponent, modulus)
        is_relative_phi = calc.relatively_prime(public_exponent, phi_modulus)
        if is_relative and is_relative_phi: break

    # retrieves the result of the extended euclid greatest common divisor,
    # this value is going to be used as the basis for the calculus of the
    # private exponent for the current operation
    d, l, _e = calc.egcd(public_exponent, phi_modulus)
    private_exponent = l

    # in case the greatest common divisor between both is not one, the values
    # are not relative primes and an exception must be raised
    if not d == 1: raise netius.GeneratorError(
        "The public exponent '%d' and the phi modulus '%d' are not relative primes" %
        (public_exponent, phi_modulus)
    )

    # calculates the inverse modulus for both exponent and in case it's not one
    # an exception is raised about the problem
    inverse_modulus = (public_exponent * private_exponent) % phi_modulus
    if not inverse_modulus == 1: netius.GeneratorError(
        "The public exponent '%d' and private exponent '%d' are not multiplicative inverse modulus of phi modulus '%d'" %
        (public_exponent, private_exponent, phi_modulus)
    )

    # creates the tuple that contains both the public and the private
    # exponent values that may be used for rsa based cryptography
    return (public_exponent, private_exponent)

def rsa_bits(modulus):
    bits = math.log(modulus, 2)
    return calc.ceil_integer(bits)

def rsa_sign(message, private_key):
    message = netius.legacy.bytes(message)
    modulus = private_key["modulus"]
    private_exponent = private_key["private_exponent"]
    return rsa_crypt_s(message, private_exponent, modulus)

def rsa_verify(signature, public_key):
    signature = netius.legacy.bytes(signature)
    modulus = public_key["modulus"]
    public_exponent = public_key["public_exponent"]
    return rsa_crypt_s(signature, public_exponent, modulus)

def rsa_crypt_s(message, exponent, modulus):
    modulus_l = calc.ceil_integer(math.log(modulus, 256))

    message_i = util.bytes_to_integer(message)
    message_crypt = rsa_crypt(message_i, exponent, modulus)
    message_crypt_s = util.integer_to_bytes(message_crypt, modulus_l)

    return message_crypt_s

def rsa_crypt(number, exponent, modulus):
    if not type(number) in netius.legacy.INTEGERS:
        raise TypeError("you must pass a long or an int")

    if number > 0 and math.floor(math.log(number, 2)) > math.floor(math.log(modulus, 2)):
        raise OverflowError("the message is too long")

    return pow(number, exponent, modulus)
<|endoftext|>