<commit_before>from __future__ import print_function

from threading import Thread
from time import sleep

from acq4.drivers.sensapex import UMP
from acq4.util import Qt
from .PressureControl import PressureControl, PressureControlWidget

class SensapexPressureControl(PressureControl):
    """Pressure control device driven by Sensapex analog/digital channels.
    """
    sigMeasuredPressureChanged = Qt.Signal(object, object)  # self, pressure

    def __init__(self, manager, config, name):
        self.devId = config.get('deviceId')
        if manager.config.get("drivers", {}).get("sensapex", {}).get("driverPath", None) is not None:
            UMP.set_library_path(manager.config["drivers"]["sensapex"]["driverPath"])
        address = config.pop('address', None)
        group = config.pop('group', None)
        ump = UMP.get_ump(address=address, group=group)
        self.dev = ump.get_device(self.devId)
        config.setdefault("maximum", 7e4)
        config.setdefault("minimum", -7e4)
        PressureControl.__init__(self, manager, config, name)

        self.pressureChannel = config.pop('pressureChannel')
        self._valveValueBySource = {"regulator": 1, "atmosphere": 0}
        self.sources = tuple(self._valveValueBySource.keys())
        self._busy = self.dev.is_busy()
        self._measurement = self.dev.measure_pressure(self.pressureChannel)
        self.source = self.getSource()
        self.pressure = self.getPressure()
        self._pollThread = Thread(target=self._poll)
        self._pollThread.daemon = True
        self._pollThread.start()

    def _poll(self):
        while True:
            self.getBusyStatus()
            self.measurePressure()
            sleep(self.regulatorSettlingTime)

    def _setPressure(self, p):
        self.dev.set_pressure(self.pressureChannel, p / 1000.)

    def getPressure(self):
        return self.dev.get_pressure(self.pressureChannel) * 1000

    def measurePressure(self):
        pressure = self.dev.measure_pressure(self.pressureChannel)
        if pressure != self._measurement:
            self._measurement = pressure
            self.sigMeasuredPressureChanged.emit(self, pressure)
        return pressure

    def getSource(self):
        if self.dev.get_valve(self.pressureChannel) == 1:
            return "regulator"
        else:
            return "atmosphere"

    def _setSource(self, source):
        self.dev.set_valve(self.pressureChannel, self._valveValueBySource[source])

    def calibrate(self):
        self.dev.calibrate_pressure(self.pressureChannel)

    def getBusyStatus(self):
        busy = self.dev.is_busy()
        if busy != self._busy:
            self._busy = busy
            self.sigBusyChanged.emit(self, busy)
        return busy

    def deviceInterface(self, win):
        return SensapexPressureControlWidget(dev=self)

class SensapexPressureControlWidget(Qt.QWidget):
    """Supports measured pressure display and calibration"""

    def __init__(self, dev):
        super(SensapexPressureControlWidget, self).__init__()
        self.dev = dev
        self.layout = Qt.QGridLayout()
        self.setLayout(self.layout)
        self.layout.setContentsMargins(0, 0, 0, 0)

        self.calibrateButton = Qt.QPushButton("Calibrate")
        self.calibrateButton.clicked.connect(self.dev.calibrate)
        self.layout.addWidget(self.calibrateButton, 0, 0)

        self.controlWidget = PressureControlWidget(self, dev)
        self.layout.addWidget(self.controlWidget, 0, 1)

        self.measurement = Qt.QLabel("-")
        self.measurement.setMargin(20)
        self.layout.addWidget(self.measurement, 0, 2)
        self._measurementChanged(dev, dev.measurePressure())
        dev.sigMeasuredPressureChanged.connect(self._measurementChanged)

        self._busyChanged(dev, dev.getBusyStatus())
        dev.sigBusyChanged.connect(self._busyChanged)

    def _measurementChanged(self, dev, pressure):
        self.measurement.setText(f"{pressure} kPa")

    def _busyChanged(self, dev, isBusy):
        self.calibrateButton.setEnabled(not isBusy)
<commit_msg>format measurement<commit_after>from __future__ import print_function

from threading import Thread
from time import sleep

from acq4.drivers.sensapex import UMP
from acq4.util import Qt
from .PressureControl import PressureControl, PressureControlWidget

class SensapexPressureControl(PressureControl):
    """Pressure control device driven by Sensapex analog/digital channels.
    """
    sigMeasuredPressureChanged = Qt.Signal(object, object)  # self, pressure

    def __init__(self, manager, config, name):
        self.devId = config.get('deviceId')
        if manager.config.get("drivers", {}).get("sensapex", {}).get("driverPath", None) is not None:
            UMP.set_library_path(manager.config["drivers"]["sensapex"]["driverPath"])
        address = config.pop('address', None)
        group = config.pop('group', None)
        ump = UMP.get_ump(address=address, group=group)
        self.dev = ump.get_device(self.devId)
        config.setdefault("maximum", 7e4)
        config.setdefault("minimum", -7e4)
        PressureControl.__init__(self, manager, config, name)

        self.pressureChannel = config.pop('pressureChannel')
        self._valveValueBySource = {"regulator": 1, "atmosphere": 0}
        self.sources = tuple(self._valveValueBySource.keys())
        self._busy = self.dev.is_busy()
        self._measurement = self.dev.measure_pressure(self.pressureChannel)
        self.source = self.getSource()
        self.pressure = self.getPressure()
        self._pollThread = Thread(target=self._poll)
        self._pollThread.daemon = True
        self._pollThread.start()

    def _poll(self):
        while True:
            self.getBusyStatus()
            self.measurePressure()
            sleep(self.regulatorSettlingTime)

    def _setPressure(self, p):
        self.dev.set_pressure(self.pressureChannel, p / 1000.)

    def getPressure(self):
        return self.dev.get_pressure(self.pressureChannel) * 1000

    def measurePressure(self):
        pressure = self.dev.measure_pressure(self.pressureChannel)
        if pressure != self._measurement:
            self._measurement = pressure
            self.sigMeasuredPressureChanged.emit(self, pressure)
        return pressure

    def getSource(self):
        if self.dev.get_valve(self.pressureChannel) == 1:
            return "regulator"
        else:
            return "atmosphere"

    def _setSource(self, source):
        self.dev.set_valve(self.pressureChannel, self._valveValueBySource[source])

    def calibrate(self):
        self.dev.calibrate_pressure(self.pressureChannel)

    def getBusyStatus(self):
        busy = self.dev.is_busy()
        if busy != self._busy:
            self._busy = busy
            self.sigBusyChanged.emit(self, busy)
        return busy

    def deviceInterface(self, win):
        return SensapexPressureControlWidget(dev=self)

class SensapexPressureControlWidget(Qt.QWidget):
    """Supports measured pressure display and calibration"""

    def __init__(self, dev):
        super(SensapexPressureControlWidget, self).__init__()
        self.dev = dev
        self.layout = Qt.QGridLayout()
        self.setLayout(self.layout)
        self.layout.setContentsMargins(0, 0, 0, 0)

        self.calibrateButton = Qt.QPushButton("Calibrate")
        self.calibrateButton.clicked.connect(self.dev.calibrate)
        self.layout.addWidget(self.calibrateButton, 0, 0)

        self.controlWidget = PressureControlWidget(self, dev)
        self.layout.addWidget(self.controlWidget, 0, 1)

        self.measurement = Qt.QLabel("-")
        self.measurement.setMargin(20)
        self.layout.addWidget(self.measurement, 0, 2)
        self._measurementChanged(dev, dev.measurePressure())
        dev.sigMeasuredPressureChanged.connect(self._measurementChanged)

        self._busyChanged(dev, dev.getBusyStatus())
        dev.sigBusyChanged.connect(self._busyChanged)

    def _measurementChanged(self, dev, pressure):
        self.measurement.setText(f"{pressure:.04f} kPa")

    def _busyChanged(self, dev, isBusy):
        self.calibrateButton.setEnabled(not isBusy)
<|endoftext|>


<commit_before>import numpy as np
import pandas as pd

from edisgo.flex_opt import q_control
from edisgo.tools.config import Config

class TestQControl:
    def test_get_q_sign_generator(self):
        # ToDo implement
        pass

    def test_get_q_sign_load(self):
        # ToDo implement
        pass

    def test_fixed_cosphi(self):

        timeindex = pd.date_range("1/1/1970", periods=2, freq="H")
        active_power_ts = pd.DataFrame(
            data={
                "comp_mv_1": [0.5, 1.5],
                "comp_mv_2": [2.5, 3.5],
                "comp_lv_1": [0.1, 0.0],
                "comp_lv_2": [0.15, 0.07],
            },
            index=timeindex,
        )
        q_sign = pd.Series(
            [-1.0, 1.0, 1.0, -1],
            index=["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"],
        )
        power_factor = pd.Series(
            [0.9, 0.95, 1.0, 0.9],
            index=["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"],
        )

        # test with q_sign as Series and power_factor as float
        reactive_power_ts = q_control.fixed_cosphi(
            active_power_ts,
            q_sign=q_sign,
            power_factor=0.9,
        )

        assert reactive_power_ts.shape == (2, 4)
        assert np.isclose(
            reactive_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values,
            active_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values * -0.484322,
        ).all()
        assert np.isclose(
            reactive_power_ts.loc[:, "comp_lv_1"].values,
            active_power_ts.loc[:, "comp_lv_1"].values * 0.484322,
        ).all()

        # test with q_sign as int and power_factor as Series
        reactive_power_ts = q_control.fixed_cosphi(
            active_power_ts,
            q_sign=1,
            power_factor=power_factor,
        )

        assert reactive_power_ts.shape == (2, 4)
        assert np.isclose(
            reactive_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values,
            active_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values * 0.484322,
        ).all()
        assert np.isclose(
            reactive_power_ts.loc[:, "comp_lv_1"].values,
            [0.0, 0.0],
        ).all()
        assert np.isclose(
            reactive_power_ts.loc[:, "comp_mv_2"].values,
            active_power_ts.loc[:, "comp_mv_2"].values * 0.328684,
        ).all()

        # test with q_sign as int and power_factor as float
        reactive_power_ts = q_control.fixed_cosphi(
            active_power_ts,
            q_sign=1,
            power_factor=0.95,
        )

        assert reactive_power_ts.shape == (2, 4)
        assert np.isclose(
            reactive_power_ts.loc[
                :, ["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"]
            ].values,
            active_power_ts.loc[
                :, ["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"]
            ].values
            * 0.328684,
        ).all()

    def test__fixed_cosphi_default_power_factor(
        self,
    ):
        df = pd.DataFrame(
            data={"voltage_level": ["mv", "lv", "lv"]},
            index=["comp_mv_1", "comp_lv_1", "comp_lv_2"],
        )
        config = Config()

        # test for component_type="generators"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="generators", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.9, 0.95, 0.95],
        ).all()

        # test for component_type="loads"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="loads", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.9, 0.95, 0.95],
        ).all()

        # test for component_type="charging_points"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="charging_points", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [1.0, 1.0, 1.0],
        ).all()

        # test for component_type="heat_pumps"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="heat_pumps", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.98, 0.98, 0.98],
        ).all()

        # test for component_type="storage_units"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="storage_units", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.9, 0.95, 0.95],
        ).all()

    def test__fixed_cosphi_default_reactive_power_sign(
        self,
    ):

        df = pd.DataFrame(
            data={"voltage_level": ["mv", "lv", "lv"]},
            index=["comp_mv_1", "comp_lv_1", "comp_lv_2"],
        )
        config = Config()

        # test for component_type="generators"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="generators", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [-1.0, -1.0, -1.0],
        ).all()

        # test for component_type="loads"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="loads", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [1.0, 1.0, 1.0],
        ).all()

        # test for component_type="charging_points"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="charging_points", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [1.0, 1.0, 1.0],
        ).all()

        # test for component_type="heat_pumps"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="heat_pumps", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [-1.0, -1.0, -1.0],
        ).all()

        # test for component_type="storage_units"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="storage_units", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [-1.0, -1.0, -1.0],
        ).all()
<commit_msg>Add tests<commit_after>import numpy as np
import pandas as pd

from edisgo.flex_opt import q_control
from edisgo.tools.config import Config

class TestQControl:
    def test_get_q_sign_generator(self):
        assert q_control.get_q_sign_generator("Inductive") == -1
        assert q_control.get_q_sign_generator("capacitive") == 1

    def test_get_q_sign_load(self):
        assert q_control.get_q_sign_load("inductive") == 1
        assert q_control.get_q_sign_load("Capacitive") == -1

    def test_fixed_cosphi(self):

        timeindex = pd.date_range("1/1/1970", periods=2, freq="H")
        active_power_ts = pd.DataFrame(
            data={
                "comp_mv_1": [0.5, 1.5],
                "comp_mv_2": [2.5, 3.5],
                "comp_lv_1": [0.1, 0.0],
                "comp_lv_2": [0.15, 0.07],
            },
            index=timeindex,
        )
        q_sign = pd.Series(
            [-1.0, 1.0, 1.0, -1],
            index=["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"],
        )
        power_factor = pd.Series(
            [0.9, 0.95, 1.0, 0.9],
            index=["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"],
        )

        # test with q_sign as Series and power_factor as float
        reactive_power_ts = q_control.fixed_cosphi(
            active_power_ts,
            q_sign=q_sign,
            power_factor=0.9,
        )

        assert reactive_power_ts.shape == (2, 4)
        assert np.isclose(
            reactive_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values,
            active_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values * -0.484322,
        ).all()
        assert np.isclose(
            reactive_power_ts.loc[:, "comp_lv_1"].values,
            active_power_ts.loc[:, "comp_lv_1"].values * 0.484322,
        ).all()

        # test with q_sign as int and power_factor as Series
        reactive_power_ts = q_control.fixed_cosphi(
            active_power_ts,
            q_sign=1,
            power_factor=power_factor,
        )

        assert reactive_power_ts.shape == (2, 4)
        assert np.isclose(
            reactive_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values,
            active_power_ts.loc[:, ["comp_mv_1", "comp_lv_2"]].values * 0.484322,
        ).all()
        assert np.isclose(
            reactive_power_ts.loc[:, "comp_lv_1"].values,
            [0.0, 0.0],
        ).all()
        assert np.isclose(
            reactive_power_ts.loc[:, "comp_mv_2"].values,
            active_power_ts.loc[:, "comp_mv_2"].values * 0.328684,
        ).all()

        # test with q_sign as int and power_factor as float
        reactive_power_ts = q_control.fixed_cosphi(
            active_power_ts,
            q_sign=1,
            power_factor=0.95,
        )

        assert reactive_power_ts.shape == (2, 4)
        assert np.isclose(
            reactive_power_ts.loc[
                :, ["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"]
            ].values,
            active_power_ts.loc[
                :, ["comp_mv_1", "comp_mv_2", "comp_lv_1", "comp_lv_2"]
            ].values
            * 0.328684,
        ).all()

    def test__fixed_cosphi_default_power_factor(
        self,
    ):
        df = pd.DataFrame(
            data={"voltage_level": ["mv", "lv", "lv"]},
            index=["comp_mv_1", "comp_lv_1", "comp_lv_2"],
        )
        config = Config()

        # test for component_type="generators"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="generators", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.9, 0.95, 0.95],
        ).all()

        # test for component_type="loads"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="loads", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.9, 0.95, 0.95],
        ).all()

        # test for component_type="charging_points"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="charging_points", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [1.0, 1.0, 1.0],
        ).all()

        # test for component_type="heat_pumps"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="heat_pumps", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.98, 0.98, 0.98],
        ).all()

        # test for component_type="storage_units"
        pf = q_control._fixed_cosphi_default_power_factor(
            comp_df=df, component_type="storage_units", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [0.9, 0.95, 0.95],
        ).all()

    def test__fixed_cosphi_default_reactive_power_sign(
        self,
    ):

        df = pd.DataFrame(
            data={"voltage_level": ["mv", "lv", "lv"]},
            index=["comp_mv_1", "comp_lv_1", "comp_lv_2"],
        )
        config = Config()

        # test for component_type="generators"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="generators", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [-1.0, -1.0, -1.0],
        ).all()

        # test for component_type="loads"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="loads", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [1.0, 1.0, 1.0],
        ).all()

        # test for component_type="charging_points"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="charging_points", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [1.0, 1.0, 1.0],
        ).all()

        # test for component_type="heat_pumps"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="heat_pumps", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [-1.0, -1.0, -1.0],
        ).all()

        # test for component_type="storage_units"
        pf = q_control._fixed_cosphi_default_reactive_power_sign(
            comp_df=df, component_type="storage_units", configs=config
        )

        assert pf.shape == (3,)
        assert np.isclose(
            pf.loc[["comp_mv_1", "comp_lv_1", "comp_lv_2"]].values,
            [-1.0, -1.0, -1.0],
        ).all()
<|endoftext|>


<commit_before>"""Examine and query coverage in sequencing experiments.

Provides estimates of coverage intervals based on callable regions
"""
import collections
import itertools
import os
import shutil
import yaml
import pybedtools
import numpy as np
import pysam
import toolz as tz

from bcbio import bam, utils
from bcbio.bam import ref, readstats
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger
from bcbio.pipeline import datadict as dd
from bcbio.provenance import do
from bcbio.pipeline import shared
from bcbio.structural import regions
from bcbio.variation import bedutils

GENOME_COV_THRESH = 0.40  # percent of genome covered for whole genome analysis
OFFTARGET_THRESH = 0.01  # percent of offtarget reads required to be capture (not amplification) based
DEPTH_THRESHOLDS = [1, 5, 10, 20, 50, 100, 250, 500, 1000, 5000, 10000, 50000]

def assign_interval(data):
    """Identify coverage based on percent of genome covered and relation to targets.

    Classifies coverage into 3 categories:
      - genome: Full genome coverage
      - regional: Regional coverage, like exome capture, with off-target reads
      - amplicon: Amplication based regional coverage without off-target reads
    """
    if not dd.get_coverage_interval(data):
        vrs = dd.get_variant_regions_merged(data)
        callable_file = dd.get_sample_callable(data)
        if vrs:
            callable_size = pybedtools.BedTool(vrs).total_coverage()
        else:
            callable_size = pybedtools.BedTool(callable_file).total_coverage()
        total_size = sum([c.size for c in ref.file_contigs(dd.get_ref_file(data), data["config"])])
        genome_cov_pct = callable_size / float(total_size)
        if genome_cov_pct > GENOME_COV_THRESH:
            cov_interval = "genome"
            offtarget_pct = 0.0
        elif not vrs:
            cov_interval = "regional"
            offtarget_pct = 0.0
        else:
            offtarget_pct = _count_offtarget(data, dd.get_align_bam(data) or dd.get_work_bam(data),
                                             vrs or callable_file, "variant_regions")
            if offtarget_pct > OFFTARGET_THRESH:
                cov_interval = "regional"
            else:
                cov_interval = "amplicon"
        logger.info("%s: Assigned coverage as '%s' with %.1f%% genome coverage and %.1f%% offtarget coverage"
                    % (dd.get_sample_name(data), cov_interval, genome_cov_pct * 100.0, offtarget_pct * 100.0))
        data["config"]["algorithm"]["coverage_interval"] = cov_interval
    return data

def _count_offtarget(data, bam_file, bed_file, target_name):
    mapped_unique = readstats.number_of_mapped_reads(data, bam_file, keep_dups=False)
    ontarget = readstats.number_of_mapped_reads(
        data, bam_file, keep_dups=False, bed_file=bed_file, target_name=target_name)
    if mapped_unique:
        return float(mapped_unique - ontarget) / mapped_unique
    else:
        return 0.0

def calculate(bam_file, data):
    """Calculate coverage in parallel using mosdepth.

    Removes duplicates and secondary reads from the counts:
    if ( b->core.flag & (BAM_FUNMAP | BAM_FSECONDARY | BAM_FQCFAIL | BAM_FDUP) ) continue;
    """
    params = {"min": dd.get_coverage_depth_min(data)}
    variant_regions = dd.get_variant_regions_merged(data)
    if not variant_regions:
        variant_regions = _create_genome_regions(data)
    # Back compatible with previous pre-mosdepth callable files
    callable_file = os.path.join(utils.safe_makedir(os.path.join(dd.get_work_dir(data), "align",
                                                                 dd.get_sample_name(data))),
                                 "%s-coverage.callable.bed" % (dd.get_sample_name(data)))
    if not utils.file_uptodate(callable_file, bam_file):
        vr_quantize = ("0:1:%s:" % (params["min"]), ["NO_COVERAGE", "LOW_COVERAGE", "CALLABLE"])
        to_calculate = [("variant_regions", variant_regions,
                         vr_quantize, None, "coverage_perbase" in dd.get_tools_on(data)),
                        ("sv_regions", bedutils.clean_file(regions.get_sv_bed(data), data, prefix="svregions-"),
                         None, None, False),
                        ("coverage", bedutils.clean_file(dd.get_coverage(data), data, prefix="cov-"),
                         None, DEPTH_THRESHOLDS, False)]
        depth_files = {}
        for target_name, region_bed, quantize, thresholds, per_base in to_calculate:
            if region_bed:
                cur_depth = {}
                depth_info = run_mosdepth(data, target_name, region_bed, quantize=quantize, thresholds=thresholds,
                                          per_base=per_base)
                for attr in ("dist", "regions", "thresholds", "per_base"):
                    val = getattr(depth_info, attr, None)
                    if val:
                        cur_depth[attr] = val
                depth_files[target_name] = cur_depth
                if target_name == "variant_regions":
                    callable_file = depth_info.quantize
    else:
        depth_files = {}
    final_callable = _subset_to_variant_regions(callable_file, variant_regions, data)
    return final_callable, depth_files

def _create_genome_regions(data):
    """Create whole genome contigs we want to process, only non-alts.

    Skips problem contigs like HLAs for downstream analysis.
    """
    work_dir = utils.safe_makedir(os.path.join(dd.get_work_dir(data), "coverage", dd.get_sample_name(data)))
    variant_regions = os.path.join(work_dir, "target-genome.bed")
    with file_transaction(data, variant_regions) as tx_variant_regions:
        with open(tx_variant_regions, "w") as out_handle:
            for c in shared.get_noalt_contigs(data):
                out_handle.write("%s\t%s\t%s\n" % (c.name, 0, c.size))
    return variant_regions

def _subset_to_variant_regions(callable_file, variant_regions, data):
    """Subset output callable file to only variant regions of interest.
    """
    out_file = "%s-vrsubset.bed" % utils.splitext_plus(callable_file)[0]
    if not utils.file_uptodate(out_file, callable_file):
        with file_transaction(data, out_file) as tx_out_file:
            pybedtools.BedTool(callable_file).intersect(variant_regions).saveas(tx_out_file)
    return out_file

def _get_cache_file(data, target_name):
    prefix = os.path.join(
        utils.safe_makedir(os.path.join(dd.get_work_dir(data), "align", dd.get_sample_name(data))),
        "%s-coverage" % (dd.get_sample_name(data)))
    cache_file = prefix + "-" + target_name + "-stats.yaml"
    return cache_file

def _read_cache(cache_file, reuse_cmp_files):
    reuse_cmp_file = [fn for fn in reuse_cmp_files if fn]
    if all(utils.file_uptodate(cache_file, fn) for fn in reuse_cmp_file):
        with open(cache_file) as in_handle:
            return yaml.safe_load(in_handle)
    return dict()

def _write_cache(cache, cache_file):
    with open(cache_file, "w") as out_handle:
        yaml.safe_dump(cache, out_handle, default_flow_style=False, allow_unicode=False)

def get_average_coverage(target_name, bed_file, data, bam_file=None):
    if not bam_file:
        bam_file = dd.get_align_bam(data) or dd.get_work_bam(data)
    cache_file = _get_cache_file(data, target_name)
    cache = _read_cache(cache_file, [bam_file, bed_file])
    if "avg_coverage" in cache:
        return int(cache["avg_coverage"])

    if bed_file:
        avg_cov = _average_bed_coverage(bed_file, target_name, data)
    else:
        avg_cov = _average_genome_coverage(data, bam_file)

    cache["avg_coverage"] = int(avg_cov)
    _write_cache(cache, cache_file)
    return int(avg_cov)

def _average_genome_coverage(data, bam_file):
    """Quickly calculate average coverage for whole genome files using indices.

    Includes all reads, with duplicates.
    """
    total = sum([c.size for c in ref.file_contigs(dd.get_ref_file(data), data["config"])])
    read_counts = sum(x.aligned for x in bam.idxstats(bam_file, data))
    with pysam.Samfile(bam_file, "rb") as pysam_bam:
        read_size = np.median(list(itertools.islice((a.query_length for a in pysam_bam.fetch()), 1e5)))
    avg_cov = float(read_counts * read_size) / total
    return avg_cov

def _average_bed_coverage(bed_file, target_name, data):
    depth_file = regions_coverage(bed_file, target_name, data)
    avg_covs = []
    total_len = 0
    with utils.open_gzipsafe(depth_file) as fh:
        for line_tokens in (l.rstrip().split() for l in fh if not l.startswith("#")):
            line_tokens = [x for x in line_tokens if x.strip()]
            start, end = map(int, line_tokens[1:3])
            size = end - start
            avg_covs.append(float(line_tokens[-1]) * size)
            total_len += size
    avg_cov = sum(avg_covs) / total_len if total_len > 0 else 0
    return avg_cov

def regions_coverage(bed_file, target_name, data):
    """Generate coverage over regions of interest using mosdepth.
    """
    ready_bed = tz.get_in(["depth", target_name, "regions"], data)
    if ready_bed:
        return ready_bed
    else:
        return run_mosdepth(data, target_name, bed_file).regions

def run_mosdepth(data, target_name, bed_file, per_base=False, quantize=None, thresholds=None):
    """Run mosdepth generating distribution, region depth and per-base depth.
    """
    MosdepthCov = collections.namedtuple("MosdepthCov", ("dist", "per_base", "regions", "quantize", "thresholds"))
    bam_file = dd.get_align_bam(data) or dd.get_work_bam(data)
    work_dir = utils.safe_makedir(os.path.join(dd.get_work_dir(data), "coverage", dd.get_sample_name(data)))
    prefix = os.path.join(work_dir, "%s-%s" % (dd.get_sample_name(data), target_name))
    out = MosdepthCov("%s.mosdepth.%s.dist.txt" % (prefix, "region" if bed_file else "global"),
                      ("%s.per-base.bed.gz" % prefix) if per_base else None,
                      ("%s.regions.bed.gz" % prefix) if bed_file else None,
                      ("%s.quantized.bed.gz" % prefix) if quantize else None,
                      ("%s.thresholds.bed.gz" % prefix) if thresholds else None)
    if not utils.file_uptodate(out.dist, bam_file):
        with file_transaction(data, out.dist) as tx_out_file:
            tx_prefix = os.path.join(os.path.dirname(tx_out_file), os.path.basename(prefix))
            num_cores = dd.get_cores(data)
            bed_arg = ("--by %s" % bed_file) if bed_file else ""
            perbase_arg = "" if per_base else "--no-per-base"
            mapq_arg = "-Q 1" if (per_base or quantize) else ""
            if quantize:
                quant_arg = "--quantize %s" % quantize[0]
                quant_export = " && ".join(["export MOSDEPTH_Q%s=%s" % (i, x) for (i, x) in enumerate(quantize[1])])
                quant_export += " && "
            else:
                quant_arg, quant_export = "", ""

            thresholds_cmdl = ("-T " + ",".join([str(t) for t in thresholds])) if out.thresholds else ""
            cmd = ("{quant_export}mosdepth -t {num_cores} -F 1804 {mapq_arg} {perbase_arg} {bed_arg} {quant_arg} "
                   "{tx_prefix} {bam_file} {thresholds_cmdl}")
            message = "Calculating coverage: %s %s" % (dd.get_sample_name(data), target_name)
            do.run(cmd.format(**locals()), message.format(**locals()))
            if out.per_base:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.per_base)), out.per_base)
            if out.regions:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.regions)), out.regions)
            if out.quantize:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.quantize)), out.quantize)
            if out.thresholds:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.thresholds)), out.thresholds)
    return out

def coverage_region_detailed_stats(target_name, bed_file, data, out_dir):
    """
    Calculate coverage at different completeness cutoff
    for region in coverage option.
    """
    if bed_file and utils.file_exists(bed_file):
        ready_depth = tz.get_in(["depth", target_name], data)
        if ready_depth:
            cov_file = ready_depth["regions"]
            dist_file = ready_depth["dist"]
            thresholds_file = ready_depth.get("thresholds")
            out_cov_file = os.path.join(out_dir, os.path.basename(cov_file))
            out_dist_file = os.path.join(out_dir, os.path.basename(dist_file))
            out_thresholds_file = os.path.join(out_dir, os.path.basename(thresholds_file)) \
                if thresholds_file and os.path.isfile(thresholds_file) else None
            if not utils.file_uptodate(out_cov_file, cov_file):
                utils.copy_plus(cov_file, out_cov_file)
                utils.copy_plus(dist_file, out_dist_file)
                utils.copy_plus(thresholds_file, out_thresholds_file) if out_thresholds_file else None
            return [out_cov_file, out_dist_file] + ([out_thresholds_file] if out_thresholds_file else [])
    return []
<commit_msg>mosdepth: back compatible with runs from pre 0.2.2<commit_after>"""Examine and query coverage in sequencing experiments.

Provides estimates of coverage intervals based on callable regions
"""
import collections
import itertools
import os
import shutil
import yaml
import pybedtools
import numpy as np
import pysam
import toolz as tz

from bcbio import bam, utils
from bcbio.bam import ref, readstats
from bcbio.distributed.transaction import file_transaction
from bcbio.log import logger
from bcbio.pipeline import datadict as dd
from bcbio.provenance import do
from bcbio.pipeline import shared
from bcbio.structural import regions
from bcbio.variation import bedutils

GENOME_COV_THRESH = 0.40  # percent of genome covered for whole genome analysis
OFFTARGET_THRESH = 0.01  # percent of offtarget reads required to be capture (not amplification) based
DEPTH_THRESHOLDS = [1, 5, 10, 20, 50, 100, 250, 500, 1000, 5000, 10000, 50000]

def assign_interval(data):
    """Identify coverage based on percent of genome covered and relation to targets.

    Classifies coverage into 3 categories:
      - genome: Full genome coverage
      - regional: Regional coverage, like exome capture, with off-target reads
      - amplicon: Amplication based regional coverage without off-target reads
    """
    if not dd.get_coverage_interval(data):
        vrs = dd.get_variant_regions_merged(data)
        callable_file = dd.get_sample_callable(data)
        if vrs:
            callable_size = pybedtools.BedTool(vrs).total_coverage()
        else:
            callable_size = pybedtools.BedTool(callable_file).total_coverage()
        total_size = sum([c.size for c in ref.file_contigs(dd.get_ref_file(data), data["config"])])
        genome_cov_pct = callable_size / float(total_size)
        if genome_cov_pct > GENOME_COV_THRESH:
            cov_interval = "genome"
            offtarget_pct = 0.0
        elif not vrs:
            cov_interval = "regional"
            offtarget_pct = 0.0
        else:
            offtarget_pct = _count_offtarget(data, dd.get_align_bam(data) or dd.get_work_bam(data),
                                             vrs or callable_file, "variant_regions")
            if offtarget_pct > OFFTARGET_THRESH:
                cov_interval = "regional"
            else:
                cov_interval = "amplicon"
        logger.info("%s: Assigned coverage as '%s' with %.1f%% genome coverage and %.1f%% offtarget coverage"
                    % (dd.get_sample_name(data), cov_interval, genome_cov_pct * 100.0, offtarget_pct * 100.0))
        data["config"]["algorithm"]["coverage_interval"] = cov_interval
    return data

def _count_offtarget(data, bam_file, bed_file, target_name):
    mapped_unique = readstats.number_of_mapped_reads(data, bam_file, keep_dups=False)
    ontarget = readstats.number_of_mapped_reads(
        data, bam_file, keep_dups=False, bed_file=bed_file, target_name=target_name)
    if mapped_unique:
        return float(mapped_unique - ontarget) / mapped_unique
    else:
        return 0.0

def calculate(bam_file, data):
    """Calculate coverage in parallel using mosdepth.

    Removes duplicates and secondary reads from the counts:
    if ( b->core.flag & (BAM_FUNMAP | BAM_FSECONDARY | BAM_FQCFAIL | BAM_FDUP) ) continue;
    """
    params = {"min": dd.get_coverage_depth_min(data)}
    variant_regions = dd.get_variant_regions_merged(data)
    if not variant_regions:
        variant_regions = _create_genome_regions(data)
    # Back compatible with previous pre-mosdepth callable files
    callable_file = os.path.join(utils.safe_makedir(os.path.join(dd.get_work_dir(data), "align",
                                                                 dd.get_sample_name(data))),
                                 "%s-coverage.callable.bed" % (dd.get_sample_name(data)))
    if not utils.file_uptodate(callable_file, bam_file):
        vr_quantize = ("0:1:%s:" % (params["min"]), ["NO_COVERAGE", "LOW_COVERAGE", "CALLABLE"])
        to_calculate = [("variant_regions", variant_regions,
                         vr_quantize, None, "coverage_perbase" in dd.get_tools_on(data)),
                        ("sv_regions", bedutils.clean_file(regions.get_sv_bed(data), data, prefix="svregions-"),
                         None, None, False),
                        ("coverage", bedutils.clean_file(dd.get_coverage(data), data, prefix="cov-"),
                         None, DEPTH_THRESHOLDS, False)]
        depth_files = {}
        for target_name, region_bed, quantize, thresholds, per_base in to_calculate:
            if region_bed:
                cur_depth = {}
                depth_info = run_mosdepth(data, target_name, region_bed, quantize=quantize, thresholds=thresholds,
                                          per_base=per_base)
                for attr in ("dist", "regions", "thresholds", "per_base"):
                    val = getattr(depth_info, attr, None)
                    if val:
                        cur_depth[attr] = val
                depth_files[target_name] = cur_depth
                if target_name == "variant_regions":
                    callable_file = depth_info.quantize
    else:
        depth_files = {}
    final_callable = _subset_to_variant_regions(callable_file, variant_regions, data)
    return final_callable, depth_files

def _create_genome_regions(data):
    """Create whole genome contigs we want to process, only non-alts.

    Skips problem contigs like HLAs for downstream analysis.
    """
    work_dir = utils.safe_makedir(os.path.join(dd.get_work_dir(data), "coverage", dd.get_sample_name(data)))
    variant_regions = os.path.join(work_dir, "target-genome.bed")
    with file_transaction(data, variant_regions) as tx_variant_regions:
        with open(tx_variant_regions, "w") as out_handle:
            for c in shared.get_noalt_contigs(data):
                out_handle.write("%s\t%s\t%s\n" % (c.name, 0, c.size))
    return variant_regions

def _subset_to_variant_regions(callable_file, variant_regions, data):
    """Subset output callable file to only variant regions of interest.
    """
    out_file = "%s-vrsubset.bed" % utils.splitext_plus(callable_file)[0]
    if not utils.file_uptodate(out_file, callable_file):
        with file_transaction(data, out_file) as tx_out_file:
            pybedtools.BedTool(callable_file).intersect(variant_regions).saveas(tx_out_file)
    return out_file

def _get_cache_file(data, target_name):
    prefix = os.path.join(
        utils.safe_makedir(os.path.join(dd.get_work_dir(data), "align", dd.get_sample_name(data))),
        "%s-coverage" % (dd.get_sample_name(data)))
    cache_file = prefix + "-" + target_name + "-stats.yaml"
    return cache_file

def _read_cache(cache_file, reuse_cmp_files):
    reuse_cmp_file = [fn for fn in reuse_cmp_files if fn]
    if all(utils.file_uptodate(cache_file, fn) for fn in reuse_cmp_file):
        with open(cache_file) as in_handle:
            return yaml.safe_load(in_handle)
    return dict()

def _write_cache(cache, cache_file):
    with open(cache_file, "w") as out_handle:
        yaml.safe_dump(cache, out_handle, default_flow_style=False, allow_unicode=False)

def get_average_coverage(target_name, bed_file, data, bam_file=None):
    if not bam_file:
        bam_file = dd.get_align_bam(data) or dd.get_work_bam(data)
    cache_file = _get_cache_file(data, target_name)
    cache = _read_cache(cache_file, [bam_file, bed_file])
    if "avg_coverage" in cache:
        return int(cache["avg_coverage"])

    if bed_file:
        avg_cov = _average_bed_coverage(bed_file, target_name, data)
    else:
        avg_cov = _average_genome_coverage(data, bam_file)

    cache["avg_coverage"] = int(avg_cov)
    _write_cache(cache, cache_file)
    return int(avg_cov)

def _average_genome_coverage(data, bam_file):
    """Quickly calculate average coverage for whole genome files using indices.

    Includes all reads, with duplicates.
    """
    total = sum([c.size for c in ref.file_contigs(dd.get_ref_file(data), data["config"])])
    read_counts = sum(x.aligned for x in bam.idxstats(bam_file, data))
    with pysam.Samfile(bam_file, "rb") as pysam_bam:
        read_size = np.median(list(itertools.islice((a.query_length for a in pysam_bam.fetch()), 1e5)))
    avg_cov = float(read_counts * read_size) / total
    return avg_cov

def _average_bed_coverage(bed_file, target_name, data):
    depth_file = regions_coverage(bed_file, target_name, data)
    avg_covs = []
    total_len = 0
    with utils.open_gzipsafe(depth_file) as fh:
        for line_tokens in (l.rstrip().split() for l in fh if not l.startswith("#")):
            line_tokens = [x for x in line_tokens if x.strip()]
            start, end = map(int, line_tokens[1:3])
            size = end - start
            avg_covs.append(float(line_tokens[-1]) * size)
            total_len += size
    avg_cov = sum(avg_covs) / total_len if total_len > 0 else 0
    return avg_cov

def regions_coverage(bed_file, target_name, data):
    """Generate coverage over regions of interest using mosdepth.
    """
    ready_bed = tz.get_in(["depth", target_name, "regions"], data)
    if ready_bed:
        return ready_bed
    else:
        return run_mosdepth(data, target_name, bed_file).regions

def run_mosdepth(data, target_name, bed_file, per_base=False, quantize=None, thresholds=None):
    """Run mosdepth generating distribution, region depth and per-base depth.
    """
    MosdepthCov = collections.namedtuple("MosdepthCov", ("dist", "per_base", "regions", "quantize", "thresholds"))
    bam_file = dd.get_align_bam(data) or dd.get_work_bam(data)
    work_dir = utils.safe_makedir(os.path.join(dd.get_work_dir(data), "coverage", dd.get_sample_name(data)))
    prefix = os.path.join(work_dir, "%s-%s" % (dd.get_sample_name(data), target_name))
    old_dist_file = "%s.mosdepth.dist.txt" % (prefix)
    out = MosdepthCov((old_dist_file if utils.file_uptodate(old_dist_file, bam_file) else
                       "%s.mosdepth.%s.dist.txt" % (prefix, "region" if bed_file else "global")),
                      ("%s.per-base.bed.gz" % prefix) if per_base else None,
                      ("%s.regions.bed.gz" % prefix) if bed_file else None,
                      ("%s.quantized.bed.gz" % prefix) if quantize else None,
                      ("%s.thresholds.bed.gz" % prefix) if thresholds else None)
    if not utils.file_uptodate(out.dist, bam_file):
        with file_transaction(data, out.dist) as tx_out_file:
            tx_prefix = os.path.join(os.path.dirname(tx_out_file), os.path.basename(prefix))
            num_cores = dd.get_cores(data)
            bed_arg = ("--by %s" % bed_file) if bed_file else ""
            perbase_arg = "" if per_base else "--no-per-base"
            mapq_arg = "-Q 1" if (per_base or quantize) else ""
            if quantize:
                quant_arg = "--quantize %s" % quantize[0]
                quant_export = " && ".join(["export MOSDEPTH_Q%s=%s" % (i, x) for (i, x) in enumerate(quantize[1])])
                quant_export += " && "
            else:
                quant_arg, quant_export = "", ""

            thresholds_cmdl = ("-T " + ",".join([str(t) for t in thresholds])) if out.thresholds else ""
            cmd = ("{quant_export}mosdepth -t {num_cores} -F 1804 {mapq_arg} {perbase_arg} {bed_arg} {quant_arg} "
                   "{tx_prefix} {bam_file} {thresholds_cmdl}")
            message = "Calculating coverage: %s %s" % (dd.get_sample_name(data), target_name)
            do.run(cmd.format(**locals()), message.format(**locals()))
            if out.per_base:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.per_base)), out.per_base)
            if out.regions:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.regions)), out.regions)
            if out.quantize:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.quantize)), out.quantize)
            if out.thresholds:
                shutil.move(os.path.join(os.path.dirname(tx_out_file), os.path.basename(out.thresholds)), out.thresholds)
    return out

def coverage_region_detailed_stats(target_name, bed_file, data, out_dir):
    """
    Calculate coverage at different completeness cutoff
    for region in coverage option.
    """
    if bed_file and utils.file_exists(bed_file):
        ready_depth = tz.get_in(["depth", target_name], data)
        if ready_depth:
            cov_file = ready_depth["regions"]
            dist_file = ready_depth["dist"]
            thresholds_file = ready_depth.get("thresholds")
            out_cov_file = os.path.join(out_dir, os.path.basename(cov_file))
            out_dist_file = os.path.join(out_dir, os.path.basename(dist_file))
            out_thresholds_file = os.path.join(out_dir, os.path.basename(thresholds_file)) \
                if thresholds_file and os.path.isfile(thresholds_file) else None
            if not utils.file_uptodate(out_cov_file, cov_file):
                utils.copy_plus(cov_file, out_cov_file)
                utils.copy_plus(dist_file, out_dist_file)
                utils.copy_plus(thresholds_file, out_thresholds_file) if out_thresholds_file else None
            return [out_cov_file, out_dist_file] + ([out_thresholds_file] if out_thresholds_file else [])
    return []
<|endoftext|>


<commit_before>"""High level parallel SNP and indel calling using multiple variant callers.
"""
import os
import collections
import copy

from bcbio import utils
from bcbio.distributed.split import grouped_parallel_split_combine
from bcbio.pipeline.shared import process_bam_by_chromosome
from bcbio.variation import gatk, gatkfilter, multi, phasing, ploidy, vfilter

# ## Variant filtration -- shared functionality

def variant_filtration(call_file, ref_file, vrn_files, data):
    """Filter variant calls using Variant Quality Score Recalibration.

    Newer GATK with Haplotype calling has combined SNP/indel filtering.
    """
    caller = data["config"]["algorithm"].get("variantcaller")
    call_file = ploidy.filter_vcf_by_sex(call_file, data)
    if caller in ["freebayes"]:
        return vfilter.freebayes(call_file, ref_file, vrn_files, data)
    elif caller in ["gatk", "gatk-haplotype"]:
        return gatkfilter.run(call_file, ref_file, vrn_files, data)
    # no additional filtration for callers that filter as part of call process
    else:
        return call_file

# ## High level functionality to run genotyping in parallel

def get_variantcaller(data):
    return data["config"]["algorithm"].get("variantcaller", "gatk")

def combine_multiple_callers(data):
    """Collapse together variant calls from multiple approaches into variants
    """
    by_bam = collections.defaultdict(list)
    for x in data:
        by_bam[x[0]["work_bam"]].append(x[0])
    out = []
    for grouped_calls in by_bam.itervalues():
        ready_calls = [{"variantcaller": get_variantcaller(x),
                        "vrn_file": x.get("vrn_file"),
                        "validate": x.get("validate")}
                       for x in grouped_calls]
        final = grouped_calls[0]
        def orig_variantcaller_order(x):
            return final["config"]["algorithm"]["orig_variantcaller"].index(x["variantcaller"])
        if len(ready_calls) > 1 and "orig_variantcaller" in final["config"]["algorithm"]:
            final["variants"] = sorted(ready_calls, key=orig_variantcaller_order)
        else:
            final["variants"] = ready_calls
        out.append([final])
    return out

def handle_multiple_variantcallers(data):
    """Split samples that potentially require multiple variant calling approaches.
    """
    assert len(data) == 1
    callers = get_variantcaller(data[0])
    if isinstance(callers, basestring):
        return [data]
    elif not callers:
        return []
    else:
        out = []
        for caller in callers:
            base = copy.deepcopy(data[0])
            base["config"]["algorithm"]["orig_variantcaller"] = \
              base["config"]["algorithm"]["variantcaller"]
            base["config"]["algorithm"]["variantcaller"] = caller
            out.append([base])
        return out

def parallel_variantcall(sample_info, parallel_fn):
    """Provide sample genotyping, running in parallel over individual chromosomes.
    """
    to_process = []
    finished = []
    for x in sample_info:
        if get_variantcaller(x[0]):
            to_process.extend(handle_multiple_variantcallers(x))
        else:
            finished.append(x)
    if len(to_process) > 0:
        split_fn = process_bam_by_chromosome("-variants.vcf.gz", "work_bam",
                                             dir_ext_fn=get_variantcaller)
        processed = grouped_parallel_split_combine(
            to_process, split_fn, multi.group_batches, parallel_fn,
            "variantcall_sample", "split_variants_by_sample", "combine_variant_files",
            "vrn_file", ["sam_ref", "config"])
        finished.extend(processed)
    return finished

def get_variantcallers():
    from bcbio.variation import freebayes, cortex, samtools, varscan, mutect
    return {"gatk": gatk.unified_genotyper,
            "gatk-haplotype": gatk.haplotype_caller,
            "freebayes": freebayes.run_freebayes,
            "cortex": cortex.run_cortex,
            "samtools": samtools.run_samtools,
            "varscan": varscan.run_varscan,
            "mutect": mutect.mutect_caller}

def variantcall_sample(data, region=None, out_file=None):
    """Parallel entry point for doing genotyping of a region of a sample.
    """
    if out_file is None or not os.path.exists(out_file) or not os.path.lexists(out_file):
        utils.safe_makedir(os.path.dirname(out_file))
        sam_ref = data["sam_ref"]
        config = data["config"]
        caller_fns = get_variantcallers()
        caller_fn = caller_fns[config["algorithm"].get("variantcaller", "gatk")]
        if isinstance(data["work_bam"], basestring):
            align_bams = [data["work_bam"]]
            items = [data]
        else:
            align_bams = data["work_bam"]
            items = data["work_items"]
        call_file = "%s-raw%s" % utils.splitext_plus(out_file)
        call_file = caller_fn(align_bams, items, sam_ref,
                              data["genome_resources"]["variation"],
                              region, call_file)
        if data["config"]["algorithm"].get("phasing", False) == "gatk":
            call_file = phasing.read_backed_phasing(call_file, align_bams, sam_ref, region, config)
        utils.symlink_plus(call_file, out_file)
        if "work_items" in data:
            del data["work_items"]
    data["vrn_file"] = out_file
    return [data]
<commit_msg>Ensure multiple callers are correctly merged for split BAMs. Thanks to @mjafin<commit_after>"""High level parallel SNP and indel calling using multiple variant callers.
"""
import os
import collections
import copy

from bcbio import utils
from bcbio.distributed.split import grouped_parallel_split_combine
from bcbio.pipeline.shared import process_bam_by_chromosome
from bcbio.variation import gatk, gatkfilter, multi, phasing, ploidy, vfilter

# ## Variant filtration -- shared functionality

def variant_filtration(call_file, ref_file, vrn_files, data):
    """Filter variant calls using Variant Quality Score Recalibration.

    Newer GATK with Haplotype calling has combined SNP/indel filtering.
    """
    caller = data["config"]["algorithm"].get("variantcaller")
    call_file = ploidy.filter_vcf_by_sex(call_file, data)
    if caller in ["freebayes"]:
        return vfilter.freebayes(call_file, ref_file, vrn_files, data)
    elif caller in ["gatk", "gatk-haplotype"]:
        return gatkfilter.run(call_file, ref_file, vrn_files, data)
    # no additional filtration for callers that filter as part of call process
    else:
        return call_file

# ## High level functionality to run genotyping in parallel

def get_variantcaller(data):
    return data["config"]["algorithm"].get("variantcaller", "gatk")

def combine_multiple_callers(data):
    """Collapse together variant calls from multiple approaches into variants
    """
    by_bam = collections.defaultdict(list)
    for x in data:
        work_bam = utils.get_in(x[0], ("combine", "work_bam", "out"), x[0]["work_bam"])
        by_bam[work_bam].append(x[0])
    out = []
    for grouped_calls in by_bam.itervalues():
        ready_calls = [{"variantcaller": get_variantcaller(x),
                        "vrn_file": x.get("vrn_file"),
                        "validate": x.get("validate")}
                       for x in grouped_calls]
        final = grouped_calls[0]
        def orig_variantcaller_order(x):
            return final["config"]["algorithm"]["orig_variantcaller"].index(x["variantcaller"])
        if len(ready_calls) > 1 and "orig_variantcaller" in final["config"]["algorithm"]:
            final["variants"] = sorted(ready_calls, key=orig_variantcaller_order)
        else:
            final["variants"] = ready_calls
        out.append([final])
    return out

def handle_multiple_variantcallers(data):
    """Split samples that potentially require multiple variant calling approaches.
    """
    assert len(data) == 1
    callers = get_variantcaller(data[0])
    if isinstance(callers, basestring):
        return [data]
    elif not callers:
        return []
    else:
        out = []
        for caller in callers:
            base = copy.deepcopy(data[0])
            base["config"]["algorithm"]["orig_variantcaller"] = \
              base["config"]["algorithm"]["variantcaller"]
            base["config"]["algorithm"]["variantcaller"] = caller
            out.append([base])
        return out

def parallel_variantcall(sample_info, parallel_fn):
    """Provide sample genotyping, running in parallel over individual chromosomes.
    """
    to_process = []
    finished = []
    for x in sample_info:
        if get_variantcaller(x[0]):
            to_process.extend(handle_multiple_variantcallers(x))
        else:
            finished.append(x)
    if len(to_process) > 0:
        split_fn = process_bam_by_chromosome("-variants.vcf.gz", "work_bam",
                                             dir_ext_fn=get_variantcaller)
        processed = grouped_parallel_split_combine(
            to_process, split_fn, multi.group_batches, parallel_fn,
            "variantcall_sample", "split_variants_by_sample", "combine_variant_files",
            "vrn_file", ["sam_ref", "config"])
        finished.extend(processed)
    return finished

def get_variantcallers():
    from bcbio.variation import freebayes, cortex, samtools, varscan, mutect
    return {"gatk": gatk.unified_genotyper,
            "gatk-haplotype": gatk.haplotype_caller,
            "freebayes": freebayes.run_freebayes,
            "cortex": cortex.run_cortex,
            "samtools": samtools.run_samtools,
            "varscan": varscan.run_varscan,
            "mutect": mutect.mutect_caller}

def variantcall_sample(data, region=None, out_file=None):
    """Parallel entry point for doing genotyping of a region of a sample.
    """
    if out_file is None or not os.path.exists(out_file) or not os.path.lexists(out_file):
        utils.safe_makedir(os.path.dirname(out_file))
        sam_ref = data["sam_ref"]
        config = data["config"]
        caller_fns = get_variantcallers()
        caller_fn = caller_fns[config["algorithm"].get("variantcaller", "gatk")]
        if isinstance(data["work_bam"], basestring):
            align_bams = [data["work_bam"]]
            items = [data]
        else:
            align_bams = data["work_bam"]
            items = data["work_items"]
        call_file = "%s-raw%s" % utils.splitext_plus(out_file)
        call_file = caller_fn(align_bams, items, sam_ref,
                              data["genome_resources"]["variation"],
                              region, call_file)
        if data["config"]["algorithm"].get("phasing", False) == "gatk":
            call_file = phasing.read_backed_phasing(call_file, align_bams, sam_ref, region, config)
        utils.symlink_plus(call_file, out_file)
        if "work_items" in data:
            del data["work_items"]
    data["vrn_file"] = out_file
    return [data]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# This file is part of beets.
# Copyright 2015-2016, Ohm Patel.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Fetch various AcousticBrainz metadata using MBID.
"""
from __future__ import division, absolute_import, print_function

import requests
import operator

from beets import plugins, ui
from functools import reduce

ACOUSTIC_BASE = "https://acousticbrainz.org/"
LEVELS = ["/low-level", "/high-level"]

class AcousticPlugin(plugins.BeetsPlugin):
    def __init__(self):
        super(AcousticPlugin, self).__init__()

        self.config.add({'auto': True})
        if self.config['auto']:
            self.register_listener('import_task_files',
                                   self.import_task_files)

    def commands(self):
        cmd = ui.Subcommand('acousticbrainz',
                            help=u"fetch metadata from AcousticBrainz")

        def func(lib, opts, args):
            items = lib.items(ui.decargs(args))
            self._fetch_info(items, ui.should_write())

        cmd.func = func
        return [cmd]

    def import_task_files(self, session, task):
        """Function is called upon beet import.
        """
        self._fetch_info(task.imported_items(), False)

    def _fetch_info(self, items, write):
        """Get data from AcousticBrainz for the items.
        """

        def get_value(*map_path):
            try:
                return reduce(operator.getitem, map_path, data)
            except KeyError:
                log.debug(u'Invalid Path: {}', map_path)

        for item in (item for item in items if item.mb_trackid):
            self._log.info(u'getting data for: {}', item)
            data = self._get_data(item.mb_trackid)
            if data:
                # Get each field and assign it on the item.
                item.danceable = get_value(
                    "highlevel", "danceability", "all", "danceable",
                )
                item.gender = get_value(
                    "highlevel", "gender", "value",
                )
                item.genre_rosamerica = get_value(
                    "highlevel", "genre_rosamerica", "value"
                )
                item.mood_acoustic = get_value(
                    "highlevel", "mood_acoustic", "all", "acoustic"
                )
                item.mood_aggressive = get_value(
                    "highlevel", "mood_aggressive", "all", "aggressive"
                )
                item.mood_electronic = get_value(
                    "highlevel", "mood_electronic", "all", "electronic"
                )
                item.mood_happy = get_value(
                    "highlevel", "mood_happy", "all", "happy"
                )
                item.mood_party = get_value(
                    "highlevel", "mood_party", "all", "party"
                )
                item.mood_relaxed = get_value(
                    "highlevel", "mood_relaxed", "all", "relaxed"
                )
                item.mood_sad = get_value(
                    "highlevel", "mood_sad", "all", "sad"
                )
                item.rhythm = get_value(
                    "highlevel", "ismir04_rhythm", "value"
                )
                item.tonal = get_value(
                    "highlevel", "tonal_atonal", "all", "tonal"
                )
                item.voice_instrumental = get_value(
                    "highlevel", "voice_instrumental", "value"
                )
                item.average_loudness = get_value(
                    "lowlevel", "average_loudness"
                )
                item.bpm = get_value(
                    "rhythm", "bpm"
                )
                item.chords_changes_rate = get_value(
                    "tonal", "chords_changes_rate"
                )
                item.chords_key = get_value(
                    "tonal", "chords_key"
                )
                item.chords_number_rate = get_value(
                    "tonal", "chords_number_rate"
                )
                item.chords_scale = get_value(
                    "tonal", "chords_scale"
                )
                item.initial_key = '{} {}'.format(
                    get_value("tonal", "key_key"),
                    get_value("tonal", "key_scale")
                )
                item.key_strength = get_value(
                    "tonal", "key_strength"
                )

                # Store the data.
                item.store()
                if write:
                    item.try_write()

    def _get_data(self, mbid):
        data = {}
        for url in _generate_urls(mbid):
            self._log.debug(u'fetching URL: {}', url)

            try:
                res = requests.get(url)
            except requests.RequestException as exc:
                self._log.info(u'request error: {}', exc)
                return {}

            if res.status_code == 404:
                self._log.info(u'recording ID {} not found', mbid)
                return {}

            try:
                data.update(res.json())
            except ValueError:
                self._log.debug(u'Invalid Response: {}', res.text)
                return {}

        return data

def _generate_urls(mbid):
    """Generates AcousticBrainz end point url for given MBID.
    """
    for level in LEVELS:
        yield ACOUSTIC_BASE + mbid + level
<commit_msg>acousticbrainz: Refactor the json handling part<commit_after># -*- coding: utf-8 -*-
# This file is part of beets.
# Copyright 2015-2016, Ohm Patel.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Fetch various AcousticBrainz metadata using MBID.
"""
from __future__ import division, absolute_import, print_function

import requests

from collections import defaultdict
from beets.util.collections import DefaultList
from beets import plugins, ui

ACOUSTIC_BASE = "https://acousticbrainz.org/"
LEVELS = ["/low-level", "/high-level"]
ABSCHEME = {
    'highlevel': {
        'danceability': {
            'all': {
                'danceable': 'danceable'
            }
        },
        'gender': {
            'value': 'gender'
        },
        'genre_rosamerica': {
            'value': 'genre_rosamerica'
        },
        'mood_acoustic': {
            'all': {
                'acoustic': 'mood_acoustic'
            }
        },
        'mood_aggressive': {
            'all': {
                'aggressive': 'mood_aggressive'
            }
        },
        'mood_electronic': {
            'all': {
                'electronic': 'mood_electronic'
            }
        },
        'mood_happy': {
            'all': {
                'happy': 'mood_happy'
            }
        },
        'mood_party': {
            'all': {
                'party': 'mood_party'
            }
        },
        'mood_relaxed': {
            'all': {
                'relaxed': 'mood_relaxed'
            }
        },
        'mood_sad': {
            'all': {
                'sad': 'mood_sad'
            }
        },
        'ismir04_rhythm': {
            'value': 'rhythm'
        },
        'tonal_atonal': {
            'all': {
                'tonal': 'tonal'
            }
        },
        'voice_instrumental': {
            'value': 'voice_instrumental'
        },
    },
    'lowlevel': {
        'average_loudness': 'average_loudness'
    },
    'rhythm': {
        'bpm': 'bpm'
    },
    'tonal': {
        'chords_changes_rate': 'chords_changes_rate',
        'chords_key': 'chords_key',
        'chords_number_rate': 'chords_number_rate',
        'chords_scale': 'chords_scale',
        'key_key': ('initial_key', 0),
        'key_scale': ('initial_key', 1),
        'key_strength': 'key_strength'

    }
}

class AcousticPlugin(plugins.BeetsPlugin):
    def __init__(self):
        super(AcousticPlugin, self).__init__()

        self.config.add({'auto': True})
        if self.config['auto']:
            self.register_listener('import_task_files',
                                   self.import_task_files)

    def commands(self):
        cmd = ui.Subcommand('acousticbrainz',
                            help=u"fetch metadata from AcousticBrainz")

        def func(lib, opts, args):
            items = lib.items(ui.decargs(args))
            self._fetch_info(items, ui.should_write())

        cmd.func = func
        return [cmd]

    def import_task_files(self, session, task):
        """Function is called upon beet import.
        """
        self._fetch_info(task.imported_items(), False)

    def _get_data(self, mbid):
        data = {}
        for url in _generate_urls(mbid):
            self._log.debug(u'fetching URL: {}', url)

            try:
                res = requests.get(url)
            except requests.RequestException as exc:
                self._log.info(u'request error: {}', exc)
                return {}

            if res.status_code == 404:
                self._log.info(u'recording ID {} not found', mbid)
                return {}

            try:
                data.update(res.json())
            except ValueError:
                self._log.debug(u'Invalid Response: {}', res.text)
                return {}

        return data

    def _fetch_info(self, items, write):
        """Get data from AcousticBrainz for the items.
        """
        for item in (item for item in items if item.mb_trackid):
            self._log.info(u'getting data for: {}', item)
            data = self._get_data(item.mb_trackid)
            if data:
                # Get each field and assign it on the item.
                for attr, val in self._map_dict_to_scheme(data, ABSCHEME):
                    self._log.debug(u'attribute {} of {} set to {}',
                                    attr,
                                    item,
                                    val)
                    setattr(item, attr, val)
                # Store the data.
                item.store()
                if write:
                    item.try_write()

    def _map_dict_to_scheme(self,
                            d,
                            s,
                            composites=defaultdict(lambda: DefaultList('')),
                            root=True):
        for k, v in s.items():
            if k in d:
                if type(v) == dict:
                    for yielded in self._map_dict_to_scheme(d[k],
                                                            v,
                                                            composites,
                                                            root=False):
                        yield yielded
                elif type(v) == tuple:
                    composites.setdefault(v[0], DefaultList(''))[v[1]] = d[k]
                else:
                    yield (v, d[k])
            else:
                self._log.debug(u'Data {} could not be mapped to scheme {} '
                                u'because key {} was not found', d, v, k)
        if root:
            for k, v in composites.items():
                yield k, ' '.join(v)

def _generate_urls(mbid):
    """Generates AcousticBrainz end point url for given MBID.
    """
    for level in LEVELS:
        yield ACOUSTIC_BASE + mbid + level
<|endoftext|>


<commit_before>import os
from zapv2 import ZAPv2 as ZAP
import time
import subprocess
from robot.api import logger
import base64
import uuid
import json
import requests
from datetime import datetime
import boto3

class RoboZap(object):
    ROBOT_LIBRARY_SCOPE = "GLOBAL"
    scan_id = ""
    def __init__(self):
        """
        ZAP Library can be imported with one argument

        Arguments:
            - ``proxy``: Proxy is required to initialize the ZAP Proxy at that location. This MUST include the port specification as well
            - ``port``: This is a portspecification that will be used across the suite

        Examples:

        | = Keyword Definition =  | = Description =  |

        | Library `|` RoboZap  | proxy | port |
        """
        self.zap = ZAP(proxies={"http": "http://127.0.0.1:8090/", "https": "http://127.0.0.1:8090/"})
        self.port = 8090

    def start_headless_zap(self):
        """
        Start OWASP ZAP without a GUI

        Examples:

        | start gui zap  | path | port |

        """
        try:
            cmd = "/zap/" + "zap.sh -daemon -dir /home/zap/.ZAP/ -config api.disablekey=true -port 8090".format(
                self.port
            )
            print(cmd)
            subprocess.Popen(cmd.split(" "), stdout=open(os.devnull, "w"))
            time.sleep(10)
        except IOError:
            print("ZAP Path is not configured correctly")

    def start_gui_zap(self):
        """
        Start OWASP ZAP with a GUI

        Examples:

        | start gui zap  | path | port |

        """
        try:
            cmd = "/zap/" + "zap.sh -config api.disablekey=true -port 8090".format(
                self.port
            )
            print(cmd)
            subprocess.Popen(cmd.split(" "), stdout=open(os.devnull, "w"))
            time.sleep(10)
        except IOError:
            print("ZAP Path is not configured correctly")

    def zap_open_url(self, url):
        """
        Invoke URLOpen with ZAP

        Examples:

        | zap open url  | target |

        """
        self.zap.urlopen(url)
        time.sleep(4)

    def zap_define_context(self):
        """
        Add Target to a context and use the context to perform all scanning/spidering operations

        Examples:

        | zap define context  | contextname  | target |

        """
        regex = "{0}.*".format("https://ifs.local-dev")
        context_id = self.zap.context.new_context(contextname="test")
        time.sleep(1)
        self.zap.context.include_in_context("test", regex=regex)
        time.sleep(5)
        return context_id

    def zap_start_spider(self):
        """
        Start ZAP Spider with ZAP's inbuilt spider mode

        Examples:

        | zap start spider  | target  | url |

        """
        try:

            spider_id = self.zap.spider.scan(url="https://ifs.local-dev", contextname="test")
            time.sleep(2)
            return spider_id
        except Exception as e:
            print((e.message))

    def zap_spider_status(self, spider_id):
        """
        Fetches the status for the spider id provided by the user
        Examples:
        | zap spider status  | spider_id |
        """
        while int(self.zap.spider.status(spider_id)) < 100:
            logger.info(
                "Spider running at {0}%".format(int(self.zap.spider.status(spider_id)))
            )
            time.sleep(10)

    def zap_start_ascan(self, policy="Default Policy"):
        """
        Initiates ZAP Active Scan on the target url and context

        Examples:

        | zap start ascan  | context  | url |

        """
        try:
            scan_id = self.zap.ascan.scan(
                contextid="test", url="https://ifs.local-dev", scanpolicyname=policy
            )
            time.sleep(2)
            return scan_id
        except Exception as e:
            print(e.message)

    def zap_scan_status(self):
        """
        Fetches the status for the spider id provided by the user

        Examples:

        | zap scan status  | scan_id |

        """
        while int(self.zap.ascan.status(scan_id)) < 100:
            logger.info(
                "Scan running at {0}%".format(int(self.zap.ascan.status(scan_id)))
            )
            time.sleep(10)

    def zap_write_to_json_file(self):
        """

        Fetches all the results from zap.core.alerts() and writes to json file.

        Examples:

        | zap write to json  | scan_id |

        """
        print("writing json to a file")
        core = self.zap.core
        print("zap core" + core)
        all_vuls = []
        for i, na in enumerate(core.alerts(baseurl="http://localhost:8090")):
            print("current iteration " + na)
            vul = {}
            vul["name"] = na["alert"]
            vul["confidence"] = na.get("confidence", "")
            if na.get("risk") == "High":
                vul["severity"] = 3
            elif na.get("risk") == "Medium":
                vul["severity"] = 2
            elif na.get("risk") == "Low":
                vul["severity"] = 1
            else:
                vul["severity"] = 0

            vul["cwe"] = na.get("cweid", 0)
            vul["uri"] = na.get("url", "")
            vul["param"] = na.get("param", "")
            vul["attack"] = na.get("attack", "")
            vul["evidence"] = na.get("evidence", "")
            message_id = na.get("messageId", "")
            message = core.message(message_id)
            if isinstance(message, dict):
                request = base64.b64encode(
                    "{0}{1}".format(message["requestHeader"], message["requestBody"])
                )
                response = base64.b64encode(
                    "{0}{1}".format(message["responseHeader"], message["responseBody"])
                )
                vul["request"] = request
                vul["response"] = response
                vul["rtt"] = int(message["rtt"])
                print("current value is " + vul)
                logger.info("log current value is " + vul)
            all_vuls.append(vul)

        print("json string is "+ all_vuls)
        logger.info("logging value " + all_vuls)
        filename = "{0}.json".format(str(uuid.uuid4()))
        with open(filename, "wb") as json_file:
            json_file.write(json.dumps(all_vuls))

        return filename

    def zap_write_to_orchy(self, report_file, secret, access, hook_uri):
        """
                Generates an XML Report and writes said report to orchestron over a webhook.

                Mandatory Fields:
                - Report_file: Absolute Path of Report File - JSON or XML
                - Token: Webhook Token
                - hook_uri: the unique URI to post the XML Report to

                Examples:

                | zap write to orchy  | report_file_path | token | hook_uri

        """
        # xml_report = self.zap.core.xmlreport()
        # with open('zap_scan.xml','w') as zaprep:
        #     zaprep.write(xml_report)
        try:
            files = {"file": open(report_file, "rb")}
            auth = {"Secret-Key": secret, "Access-Key": access}
            r = requests.post(hook_uri, headers=auth, files=files)
            if r.status_code == 200:
                return "Successfully posted to Orchestron"
            else:
                raise Exception("Unable to post successfully")
        except Exception as e:
            print(e)

    def zap_export_report(
        self
    ):
        """
        This functionality works on ZAP 2.7.0 only. It leverages the Export Report Library to generate a report.
        Currently ExportReport doesnt have an API endpoint in python. We will be using the default ZAP REST API for this

        :param export_file: location to which the export needs to happen. Absolute path with the export file name and extension
        :param export_format: file extension of the exported file. Can be XML, JSON, HTML, PDF, DOC
        :param report_title: Title of the exported report
        :param report_author: Name of the Author of the report
        Examples:

        | zap export report | export_path | export_format |

        """

        url = "http://localhost:8090/JSON/exportreport/action/generate/".format(
            self.port
        )
        export_path = "/zap/zapreport.json"
        extension = json
        report_time = datetime.now().strftime("%I:%M%p on %B %d, %Y")
        source_info = "{0};{1};ZAP Team;{2};{3};v1;v1;{4}".format(
            "ZAP report", "Arun", report_time, report_time, "ZAP report"
        )
        alert_severity = "t;t;t;t"  # High;Medium;Low;Info
        alert_details = "t;t;t;t;t;t;t;t;t;t"  # CWEID;#WASCID;Description;Other Info;Solution;Reference;Request Header;Response Header;Request Body;Response Body
        data = {
            "absolutePath": export_path,
            "fileExtension": extension,
            "sourceDetails": source_info,
            "alertSeverity": alert_severity,
            "alertDetails": alert_details,
        }

        r = requests.post(url, data=data)
        if r.status_code == 200:
            pass
        else:
            raise Exception("Unable to generate report")

    def zap_write_to_s3_bucket(self, filename, bucket_name):
        s3 = boto3.client("s3")
        outfile_name = "ZAP-RESULT-{}.json".format(str(uuid.uuid4()))
        s3.upload_file(filename, bucket_name, outfile_name)
        logger.warn("Filename uploaded to S3 is: {}".format(outfile_name))

    def retrieve_secret_from_ssm(self, secret, region="us-west-2", decrypt=True):
        db = boto3.client("ssm", region_name=region)
        param = db.get_parameter(Name=secret, WithDecryption=decrypt)["Parameter"][
            "Value"
        ]
        return param

    def zap_load_script(
        self,
        script_name,
        script_type,
        script_engine,
        script_file,
        desc="Generic Description of a ZAP Script",
    ):
        """
        :param script_name:
        :param script_type:
        :param script_engine:
        :param script_file:
        :param desc:
        :return:
        """
        zap_script_status = self.zap.script.load(
            scriptname=script_name,
            scripttype=script_type,
            scriptengine=script_engine,
            filename=script_file,
            scriptdescription=desc,
        )
        logger.info(zap_script_status)

    def zap_run_standalone_script(self, script_name):
        zap_script_run_status = self.zap.script.run_stand_alone_script(script_name)
        logger.info(zap_script_run_status)

    def zap_shutdown(self):
        """
        Shutdown process for ZAP Scanner
        """
        self.zap.core.shutdown()<commit_msg>IFS-9451 removed print statement<commit_after>import os
from zapv2 import ZAPv2 as ZAP
import time
import subprocess
from robot.api import logger
import base64
import uuid
import json
import requests
from datetime import datetime
import boto3

class RoboZap(object):
    ROBOT_LIBRARY_SCOPE = "GLOBAL"
    scan_id = ""
    def __init__(self):
        """
        ZAP Library can be imported with one argument

        Arguments:
            - ``proxy``: Proxy is required to initialize the ZAP Proxy at that location. This MUST include the port specification as well
            - ``port``: This is a portspecification that will be used across the suite

        Examples:

        | = Keyword Definition =  | = Description =  |

        | Library `|` RoboZap  | proxy | port |
        """
        self.zap = ZAP(proxies={"http": "http://127.0.0.1:8090/", "https": "http://127.0.0.1:8090/"})
        self.port = 8090

    def start_headless_zap(self):
        """
        Start OWASP ZAP without a GUI

        Examples:

        | start gui zap  | path | port |

        """
        try:
            cmd = "/zap/" + "zap.sh -daemon -dir /home/zap/.ZAP/ -config api.disablekey=true -port 8090".format(
                self.port
            )
            print(cmd)
            subprocess.Popen(cmd.split(" "), stdout=open(os.devnull, "w"))
            time.sleep(10)
        except IOError:
            print("ZAP Path is not configured correctly")

    def start_gui_zap(self):
        """
        Start OWASP ZAP with a GUI

        Examples:

        | start gui zap  | path | port |

        """
        try:
            cmd = "/zap/" + "zap.sh -config api.disablekey=true -port 8090".format(
                self.port
            )
            print(cmd)
            subprocess.Popen(cmd.split(" "), stdout=open(os.devnull, "w"))
            time.sleep(10)
        except IOError:
            print("ZAP Path is not configured correctly")

    def zap_open_url(self, url):
        """
        Invoke URLOpen with ZAP

        Examples:

        | zap open url  | target |

        """
        self.zap.urlopen(url)
        time.sleep(4)

    def zap_define_context(self):
        """
        Add Target to a context and use the context to perform all scanning/spidering operations

        Examples:

        | zap define context  | contextname  | target |

        """
        regex = "{0}.*".format("https://ifs.local-dev")
        context_id = self.zap.context.new_context(contextname="test")
        time.sleep(1)
        self.zap.context.include_in_context("test", regex=regex)
        time.sleep(5)
        return context_id

    def zap_start_spider(self):
        """
        Start ZAP Spider with ZAP's inbuilt spider mode

        Examples:

        | zap start spider  | target  | url |

        """
        try:

            spider_id = self.zap.spider.scan(url="https://ifs.local-dev", contextname="test")
            time.sleep(2)
            return spider_id
        except Exception as e:
            print((e.message))

    def zap_spider_status(self, spider_id):
        """
        Fetches the status for the spider id provided by the user
        Examples:
        | zap spider status  | spider_id |
        """
        while int(self.zap.spider.status(spider_id)) < 100:
            logger.info(
                "Spider running at {0}%".format(int(self.zap.spider.status(spider_id)))
            )
            time.sleep(10)

    def zap_start_ascan(self, policy="Default Policy"):
        """
        Initiates ZAP Active Scan on the target url and context

        Examples:

        | zap start ascan  | context  | url |

        """
        try:
            scan_id = self.zap.ascan.scan(
                contextid="test", url="https://ifs.local-dev", scanpolicyname=policy
            )
            time.sleep(2)
            return scan_id
        except Exception as e:
            print(e.message)

    def zap_scan_status(self):
        """
        Fetches the status for the spider id provided by the user

        Examples:

        | zap scan status  | scan_id |

        """
        while int(self.zap.ascan.status(scan_id)) < 100:
            logger.info(
                "Scan running at {0}%".format(int(self.zap.ascan.status(scan_id)))
            )
            time.sleep(10)

    def zap_write_to_json_file(self):
        """

        Fetches all the results from zap.core.alerts() and writes to json file.

        Examples:

        | zap write to json  | scan_id |

        """
        print("writing json to a file")
        core = self.zap.core

        all_vuls = []
        for i, na in enumerate(core.alerts(baseurl="http://localhost:8090")):
            print("current iteration ")
            vul = {}
            vul["name"] = na["alert"]
            vul["confidence"] = na.get("confidence", "")
            if na.get("risk") == "High":
                vul["severity"] = 3
            elif na.get("risk") == "Medium":
                vul["severity"] = 2
            elif na.get("risk") == "Low":
                vul["severity"] = 1
            else:
                vul["severity"] = 0

            vul["cwe"] = na.get("cweid", 0)
            vul["uri"] = na.get("url", "")
            vul["param"] = na.get("param", "")
            vul["attack"] = na.get("attack", "")
            vul["evidence"] = na.get("evidence", "")
            message_id = na.get("messageId", "")
            message = core.message(message_id)
            if isinstance(message, dict):
                request = base64.b64encode(
                    "{0}{1}".format(message["requestHeader"], message["requestBody"])
                )
                response = base64.b64encode(
                    "{0}{1}".format(message["responseHeader"], message["responseBody"])
                )
                vul["request"] = request
                vul["response"] = response
                vul["rtt"] = int(message["rtt"])

            all_vuls.append(vul)

        print("json string is ")

        filename = "{0}.json".format(str(uuid.uuid4()))
        with open(filename, "wb") as json_file:
            json_file.write(json.dumps(all_vuls))

        return filename

    def zap_write_to_orchy(self, report_file, secret, access, hook_uri):
        """
                Generates an XML Report and writes said report to orchestron over a webhook.

                Mandatory Fields:
                - Report_file: Absolute Path of Report File - JSON or XML
                - Token: Webhook Token
                - hook_uri: the unique URI to post the XML Report to

                Examples:

                | zap write to orchy  | report_file_path | token | hook_uri

        """
        # xml_report = self.zap.core.xmlreport()
        # with open('zap_scan.xml','w') as zaprep:
        #     zaprep.write(xml_report)
        try:
            files = {"file": open(report_file, "rb")}
            auth = {"Secret-Key": secret, "Access-Key": access}
            r = requests.post(hook_uri, headers=auth, files=files)
            if r.status_code == 200:
                return "Successfully posted to Orchestron"
            else:
                raise Exception("Unable to post successfully")
        except Exception as e:
            print(e)

    def zap_export_report(
        self
    ):
        """
        This functionality works on ZAP 2.7.0 only. It leverages the Export Report Library to generate a report.
        Currently ExportReport doesnt have an API endpoint in python. We will be using the default ZAP REST API for this

        :param export_file: location to which the export needs to happen. Absolute path with the export file name and extension
        :param export_format: file extension of the exported file. Can be XML, JSON, HTML, PDF, DOC
        :param report_title: Title of the exported report
        :param report_author: Name of the Author of the report
        Examples:

        | zap export report | export_path | export_format |

        """

        url = "http://localhost:8090/JSON/exportreport/action/generate/".format(
            self.port
        )
        export_path = "/zap/zapreport.json"
        extension = json
        report_time = datetime.now().strftime("%I:%M%p on %B %d, %Y")
        source_info = "{0};{1};ZAP Team;{2};{3};v1;v1;{4}".format(
            "ZAP report", "Arun", report_time, report_time, "ZAP report"
        )
        alert_severity = "t;t;t;t"  # High;Medium;Low;Info
        alert_details = "t;t;t;t;t;t;t;t;t;t"  # CWEID;#WASCID;Description;Other Info;Solution;Reference;Request Header;Response Header;Request Body;Response Body
        data = {
            "absolutePath": export_path,
            "fileExtension": extension,
            "sourceDetails": source_info,
            "alertSeverity": alert_severity,
            "alertDetails": alert_details,
        }

        r = requests.post(url, data=data)
        if r.status_code == 200:
            pass
        else:
            raise Exception("Unable to generate report")

    def zap_write_to_s3_bucket(self, filename, bucket_name):
        s3 = boto3.client("s3")
        outfile_name = "ZAP-RESULT-{}.json".format(str(uuid.uuid4()))
        s3.upload_file(filename, bucket_name, outfile_name)
        logger.warn("Filename uploaded to S3 is: {}".format(outfile_name))

    def retrieve_secret_from_ssm(self, secret, region="us-west-2", decrypt=True):
        db = boto3.client("ssm", region_name=region)
        param = db.get_parameter(Name=secret, WithDecryption=decrypt)["Parameter"][
            "Value"
        ]
        return param

    def zap_load_script(
        self,
        script_name,
        script_type,
        script_engine,
        script_file,
        desc="Generic Description of a ZAP Script",
    ):
        """
        :param script_name:
        :param script_type:
        :param script_engine:
        :param script_file:
        :param desc:
        :return:
        """
        zap_script_status = self.zap.script.load(
            scriptname=script_name,
            scripttype=script_type,
            scriptengine=script_engine,
            filename=script_file,
            scriptdescription=desc,
        )
        logger.info(zap_script_status)

    def zap_run_standalone_script(self, script_name):
        zap_script_run_status = self.zap.script.run_stand_alone_script(script_name)
        logger.info(zap_script_run_status)

    def zap_shutdown(self):
        """
        Shutdown process for ZAP Scanner
        """
        self.zap.core.shutdown()<|endoftext|>