<commit_before># Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""
Generate a file that sets the specified board's binhosts to include all of the
other compatible boards in this buildroot.
"""

import collections
import glob
import optparse
import os
import sys

from chromite.buildbot import cbuildbot_config
from chromite.lib import cros_build_lib

def FindCandidateBoards(board):
  """Find candidate local boards to grab prebuilts from."""
  board_no_variant = board.partition("_")[0]
  portageq_prefix = "/usr/local/bin/portageq-"
  for path in sorted(glob.glob("%s*" % portageq_prefix)):
    # Strip off the portageq prefix, leaving only the board.
    other_board = path.replace(portageq_prefix, "")

    # It is only safe to inherit prebuilts from generic boards, or from the
    # same board without the variant. This rule helps keep inheritance trees
    # sane.
    if (other_board in cbuildbot_config.generic_boards or
        other_board == board_no_variant or other_board == board):
      yield other_board

def SummarizeCompatibility(board):
  """Returns a string that will be the same for compatible boards."""
  cmd = ["portageq-%s" % board, "envvar", "ARCH", "CFLAGS"]
  return cros_build_lib.RunCommand(cmd, redirect_stdout=True,
                                   print_cmd=False).output.rstrip()

def GenerateBinhostLine(build_root, compatible_boards):
  """Generate a binhost line pulling binaries from the specified boards."""
  local_binhosts = " ".join([
      "file://localhost" + os.path.join(build_root, x, "packages")
      for x in sorted(compatible_boards)])
  return "LOCAL_BINHOST='%s'" % local_binhosts

def main(argv):
  parser = optparse.OptionParser(usage="USAGE: ./%prog --board=board [options]")

  parser.add_option("--build_root", default="/build",
                    dest="build_root",
                    help="Location of boards (normally /build)")
  parser.add_option("--board", default=None,
                    dest="board",
                    help="Board name (required).")

  flags, remaining_arguments = parser.parse_args(argv)

  if remaining_arguments or not flags.board:
    parser.print_help()
    sys.exit(1)

  by_compatibility = collections.defaultdict(set)
  compatible_boards = None
  for other_board in FindCandidateBoards(flags.board):
    compat_id = SummarizeCompatibility(other_board)
    if other_board == flags.board:
      compatible_boards = by_compatibility[compat_id]
    else:
      by_compatibility[compat_id].add(other_board)

  if compatible_boards is None:
    print >> sys.stderr, "Missing portageq wrapper for %s" % flags.board
    sys.exit(1)

  # If a non-generic prebuilt is available, ignore the generic ones,
  # since the non-generic ones will be better matches.
  compatible_specific_boards = \
      compatible_boards - set(cbuildbot_config.generic_boards)
  if compatible_specific_boards:
    compatible_boards = compatible_specific_boards

  print "# Generated by cros_generate_local_binhosts."
  print GenerateBinhostLine(flags.build_root, compatible_boards)
<commit_msg>Allow sharing of local binaries between unrelated boards.<commit_after># Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Script for calculating compatible binhosts.

Generates a file that sets the specified board's binhosts to include all of the
other compatible boards in this buildroot.
"""

import collections
import glob
import optparse
import os
import sys

from chromite.buildbot import cbuildbot_config
from chromite.lib import cros_build_lib

def FindCandidateBoards():
  """Find candidate local boards to grab prebuilts from."""
  portageq_prefix = "/usr/local/bin/portageq-"
  for path in sorted(glob.glob("%s*" % portageq_prefix)):
    # Strip off the portageq prefix, leaving only the board.
    yield path.replace(portageq_prefix, "")

def SummarizeCompatibility(board):
  """Returns a string that will be the same for compatible boards."""
  cmd = ["portageq-%s" % board, "envvar", "ARCH", "CFLAGS"]
  return cros_build_lib.RunCommand(cmd, redirect_stdout=True,
                                   print_cmd=False).output.rstrip()

def GenerateBinhostLine(build_root, compatible_boards):
  """Generate a binhost line pulling binaries from the specified boards."""
  local_binhosts = " ".join([
      "file://localhost" + os.path.join(build_root, x, "packages")
      for x in sorted(compatible_boards)])
  return "LOCAL_BINHOST='%s'" % local_binhosts

def main(argv):
  parser = optparse.OptionParser(usage="USAGE: ./%prog --board=board [options]")

  parser.add_option("--build_root", default="/build",
                    dest="build_root",
                    help="Location of boards (normally /build)")
  parser.add_option("--board", default=None,
                    dest="board",
                    help="Board name (required).")

  flags, remaining_arguments = parser.parse_args(argv)

  if remaining_arguments or not flags.board:
    parser.print_help()
    sys.exit(1)

  by_compatibility = collections.defaultdict(set)
  compatible_boards = None
  for other_board in FindCandidateBoards():
    compat_id = SummarizeCompatibility(other_board)
    if other_board == flags.board:
      compatible_boards = by_compatibility[compat_id]
    else:
      by_compatibility[compat_id].add(other_board)

  if compatible_boards is None:
    print >> sys.stderr, "Missing portageq wrapper for %s" % flags.board
    sys.exit(1)

  # If a non-generic prebuilt is available, ignore the generic ones,
  # since the non-generic ones will be better matches.
  compatible_specific_boards = \
      compatible_boards - set(cbuildbot_config.generic_boards)
  if compatible_specific_boards:
    compatible_boards = compatible_specific_boards

  print "# Generated by cros_generate_local_binhosts."
  print GenerateBinhostLine(flags.build_root, compatible_boards)
<|endoftext|>


<commit_before>"""DNSBL related tests."""

from unittest import mock

import dns.resolver
from testfixtures import LogCapture

from django.core import mail, management
from django.test import override_settings
from django.urls import reverse
from django.utils.translation import ugettext as _

from modoboa.core import factories as core_factories
from modoboa.lib.tests import ModoTestCase
from . import utils
from .. import factories, models
from ..lib import get_domain_mx_list

class MXTestCase(ModoTestCase):
    """TestCase for DNSBL related features."""

    @classmethod
    def setUpTestData(cls):  # NOQA:N802
        """Create some data."""
        super(MXTestCase, cls).setUpTestData()
        cls.domain = factories.DomainFactory(name="modoboa.org")
        # should not exist
        cls.bad_domain = factories.DomainFactory(
            name="does-not-exist.example.com")
        # Add domain admin with mailbox
        mb = factories.MailboxFactory(
            address="admin", domain=cls.bad_domain,
            user__username="admin@does-not-exist.example.com",
            user__groups=("DomainAdmins", )
        )
        cls.bad_domain.add_admin(mb.user)
        # Add domain admin with no mailbox
        admin = core_factories.UserFactory(
            username="admin2@does-not-exist.example.com",
            groups=("DomainAdmins", ))
        cls.bad_domain.add_admin(admin)

        cls.localconfig.parameters.set_value(
            "valid_mxs", "192.0.2.1 2001:db8::1")
        cls.localconfig.save()
        models.MXRecord.objects.all().delete()

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        self.assertEqual(models.MXRecord.objects.count(), 0)
        management.call_command("modo", "check_mx", "--no-dnsbl", "--ttl=0")
        self.assertTrue(
            models.MXRecord.objects.filter(domain=self.domain).exists())

        # we passed a ttl to 0. this will reset the cache this time
        qs = models.MXRecord.objects.filter(domain=self.domain)
        id_ = qs[0].id
        management.call_command("modo", "check_mx", "--no-dnsbl", "--ttl=7200")
        qs = models.MXRecord.objects.filter(domain=self.domain)
        self.assertNotEqual(id_, qs[0].id)
        id_ = qs[0].id

        # assume that mxrecords ids are the same. means that we taking care of
        # ttl
        management.call_command("modo", "check_mx", "--no-dnsbl")
        qs = models.MXRecord.objects.filter(domain=self.domain)
        self.assertEqual(id_, qs[0].id)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_single_domain_update(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Update only one domain."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        management.call_command(
            "modo", "check_mx", "--domain", self.domain.name)
        self.assertTrue(
            models.MXRecord.objects.filter(domain=self.domain).exists())
        self.assertFalse(
            models.MXRecord.objects.filter(domain=self.bad_domain).exists())

        management.call_command(
            "modo", "check_mx", "--domain", str(self.bad_domain.pk))
        self.assertFalse(
            models.MXRecord.objects.filter(domain=self.bad_domain).exists())
        self.assertEqual(len(mail.outbox), 1)

        management.call_command(
            "modo", "check_mx", "--domain", "toto.com")

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_invalid_mx(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Test to check if invalid MX records are detected."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        domain = factories.DomainFactory(name="invalid-mx.com")
        # Add domain admin with mailbox
        mb = factories.MailboxFactory(
            address="admin", domain=domain,
            user__username="admin@{}".format(domain.name),
            user__groups=("DomainAdmins", )
        )
        domain.add_admin(mb.user)
        management.call_command(
            "modo", "check_mx", "--domain", domain.name, "--no-dnsbl")
        self.assertEqual(domain.alarms.count(), 1)
        management.call_command(
            "modo", "check_mx", "--domain", domain.name, "--no-dnsbl")
        self.assertEqual(domain.alarms.count(), 1)
        self.assertEqual(len(mail.outbox), 1)

    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_get_mx_list_dns_server(self, mock_query, mock_getaddrinfo):
        """Test to get mx list from specific DNS server."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        self.set_global_parameter("custom_dns_server", "123.45.67.89")
        get_domain_mx_list("does-not-exist.example.com")

    @mock.patch("ipaddress.ip_address")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_get_domain_mx_list_logging(self, mock_query, mock_ip_address):
        """Test to get error loggins from specific DNS server."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_ip_address.side_effect = utils.mock_ip_address_result
        with LogCapture("modoboa.admin") as log:
            get_domain_mx_list("does-not-exist.example.com")
            get_domain_mx_list("no-mx.example.com")
            get_domain_mx_list("no-ns-servers.example.com")
            get_domain_mx_list("timeout.example.com")
            get_domain_mx_list("no-lookup.example.com")
            get_domain_mx_list("no-answer.example.Com")
            get_domain_mx_list("bad-response.example.com")
        log.check(
            ("modoboa.admin", "ERROR",
                _("No DNS record found for %s")
                % "does-not-exist.example.com"),
            ("modoboa.admin", "ERROR",
                _("No MX record for %s") % "no-mx.example.com"),
            ("modoboa.admin", "ERROR", _("No working name servers found")),
            ("modoboa.admin", "WARNING",
                _("DNS resolution timeout, unable to query %s at the moment")
             % "timeout.example.com"),
            ("modoboa.admin", "ERROR",
             _("No DNS record found for %s")
             % "does-not-exist.example.com"),
            ("modoboa.admin", "ERROR",
             _("No DNS record found for %s")
             % "does-not-exist.example.com"),
            ("modoboa.admin", "WARNING",
             _("Invalid IP address format for %(domain)s; %(addr)s")
             % {"domain": "bad-response.example.com", "addr": "000.0.0.0"}),
            ("modoboa.admin", "WARNING",
             _("Invalid IP address format for %(domain)s; %(addr)s")
             % {"domain": "bad-response.example.com", "addr": "000.0.0.0"}),
        )

@override_settings(DNSBL_PROVIDERS=["zen.spamhaus.org"])
class DNSBLTestCase(ModoTestCase):
    """TestCase for DNSBL related features."""

    @classmethod
    def setUpTestData(cls):  # NOQA:N802
        """Create some data."""
        super().setUpTestData()
        cls.domain = factories.DomainFactory(name="modoboa.org")
        factories.DomainFactory(name="does-not-exist.example.com")
        cls.domain2 = factories.DomainFactory(
            name="test.localhost")  # Should not be checked
        cls.domain3 = factories.DomainFactory(
            name="modoboa.com", enable_dns_checks=False)
        models.DNSBLResult.objects.all().delete()

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        self.assertEqual(models.DNSBLResult.objects.count(), 0)
        management.call_command("modo", "check_mx")
        self.assertTrue(
            models.DNSBLResult.objects.filter(domain=self.domain).exists())
        self.assertFalse(
            models.DNSBLResult.objects.filter(domain=self.domain3).exists())
        response = self.client.get(
            reverse("admin:dnsbl_domain_detail", args=[self.domain.pk]))
        self.assertEqual(response.status_code, 200)
        self.assertFalse(self.domain.uses_a_reserved_tld)
        self.assertTrue(self.domain2.uses_a_reserved_tld)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_notifications(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check notifications."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "127.0.0.4"
        management.call_command(
            "modo", "check_mx", "--email", "user@example.test")
        self.assertEqual(len(mail.outbox), 2)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command_no_dnsbl(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine without dnsbl."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        self.assertEqual(models.DNSBLResult.objects.count(), 0)
        management.call_command("modo", "check_mx", "--no-dnsbl")
        self.assertFalse(
            models.DNSBLResult.objects.filter(domain=self.domain).exists())
        response = self.client.get(
            reverse("admin:dnsbl_domain_detail", args=[self.domain.pk]))
        self.assertEqual(response.status_code, 200)

class DNSChecksTestCase(ModoTestCase):
    """A test case for DNS checks."""

    @classmethod
    def setUpTestData(cls):  # NOQA:N802
        """Create some data."""
        super(DNSChecksTestCase, cls).setUpTestData()
        cls.domain = factories.DomainFactory(name="dns-checks.com")

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"

        self.domain.enable_dkim = True
        self.domain.dkim_public_key = "XXXXX"
        self.domain.save(update_fields=["enable_dkim", "dkim_public_key"])

        management.call_command("modo", "check_mx", "--no-dnsbl", "--ttl=0")

        self.assertIsNot(self.domain.spf_record, None)
        self.assertIsNot(self.domain.dkim_record, None)
        self.assertIsNot(self.domain.dmarc_record, None)
        self.assertIsNot(self.domain.autoconfig_record, None)
        self.assertIsNot(self.domain.autodiscover_record, None)
<commit_msg>make codecov happy<commit_after>"""DNSBL related tests."""

from unittest import mock

import dns.resolver
from testfixtures import LogCapture

from django.core import mail, management
from django.test import override_settings
from django.urls import reverse
from django.utils.translation import ugettext as _

from modoboa.core import factories as core_factories
from modoboa.lib.tests import ModoTestCase
from . import utils
from .. import factories, models
from ..lib import get_domain_mx_list

class MXTestCase(ModoTestCase):
    """TestCase for DNSBL related features."""

    @classmethod
    def setUpTestData(cls):  # NOQA:N802
        """Create some data."""
        super(MXTestCase, cls).setUpTestData()
        cls.domain = factories.DomainFactory(name="modoboa.org")
        # should not exist
        cls.bad_domain = factories.DomainFactory(
            name="does-not-exist.example.com")
        # Add domain admin with mailbox
        mb = factories.MailboxFactory(
            address="admin", domain=cls.bad_domain,
            user__username="admin@does-not-exist.example.com",
            user__groups=("DomainAdmins", )
        )
        cls.bad_domain.add_admin(mb.user)
        # Add domain admin with no mailbox
        admin = core_factories.UserFactory(
            username="admin2@does-not-exist.example.com",
            groups=("DomainAdmins", ))
        cls.bad_domain.add_admin(admin)

        cls.localconfig.parameters.set_value(
            "valid_mxs", "192.0.2.1 2001:db8::1")
        cls.localconfig.save()
        models.MXRecord.objects.all().delete()

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        self.assertEqual(models.MXRecord.objects.count(), 0)
        management.call_command("modo", "check_mx", "--no-dnsbl", "--ttl=0")
        self.assertTrue(
            models.MXRecord.objects.filter(domain=self.domain).exists())

        # we passed a ttl to 0. this will reset the cache this time
        qs = models.MXRecord.objects.filter(domain=self.domain)
        id_ = qs[0].id
        management.call_command("modo", "check_mx", "--no-dnsbl", "--ttl=7200")
        qs = models.MXRecord.objects.filter(domain=self.domain)
        self.assertNotEqual(id_, qs[0].id)
        id_ = qs[0].id

        # assume that mxrecords ids are the same. means that we taking care of
        # ttl
        management.call_command("modo", "check_mx", "--no-dnsbl")
        qs = models.MXRecord.objects.filter(domain=self.domain)
        self.assertEqual(id_, qs[0].id)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_single_domain_update(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Update only one domain."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        management.call_command(
            "modo", "check_mx", "--domain", self.domain.name)
        self.assertTrue(
            models.MXRecord.objects.filter(domain=self.domain).exists())
        self.assertFalse(
            models.MXRecord.objects.filter(domain=self.bad_domain).exists())

        management.call_command(
            "modo", "check_mx", "--domain", str(self.bad_domain.pk))
        self.assertFalse(
            models.MXRecord.objects.filter(domain=self.bad_domain).exists())
        self.assertEqual(len(mail.outbox), 1)

        management.call_command(
            "modo", "check_mx", "--domain", "toto.com")

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_invalid_mx(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Test to check if invalid MX records are detected."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        domain = factories.DomainFactory(name="invalid-mx.com")
        # Add domain admin with mailbox
        mb = factories.MailboxFactory(
            address="admin", domain=domain,
            user__username="admin@{}".format(domain.name),
            user__groups=("DomainAdmins", )
        )
        domain.add_admin(mb.user)
        management.call_command(
            "modo", "check_mx", "--domain", domain.name, "--no-dnsbl")
        self.assertEqual(domain.alarms.count(), 1)
        management.call_command(
            "modo", "check_mx", "--domain", domain.name, "--no-dnsbl")
        self.assertEqual(domain.alarms.count(), 1)
        self.assertEqual(len(mail.outbox), 1)

    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_get_mx_list_dns_server(self, mock_query, mock_getaddrinfo):
        """Test to get mx list from specific DNS server."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        self.set_global_parameter("custom_dns_server", "123.45.67.89")
        get_domain_mx_list("does-not-exist.example.com")

    @mock.patch("ipaddress.ip_address")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_get_domain_mx_list_logging(self, mock_query, mock_ip_address):
        """Test to get error loggins from specific DNS server."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_ip_address.side_effect = utils.mock_ip_address_result
        with LogCapture("modoboa.admin") as log:
            get_domain_mx_list("does-not-exist.example.com")
            get_domain_mx_list("no-mx.example.com")
            get_domain_mx_list("no-ns-servers.example.com")
            get_domain_mx_list("timeout.example.com")
            get_domain_mx_list("no-lookup.example.com")
            get_domain_mx_list("no-answer.example.Com")
            get_domain_mx_list("bad-response.example.com")
        log.check(
            ("modoboa.admin", "ERROR",
                _("No DNS record found for %s")
                % "does-not-exist.example.com"),
            ("modoboa.admin", "ERROR",
                _("No MX record for %s") % "no-mx.example.com"),
            ("modoboa.admin", "ERROR", _("No working name servers found")),
            ("modoboa.admin", "WARNING",
                _("DNS resolution timeout, unable to query %s at the moment")
             % "timeout.example.com"),
            ("modoboa.admin", "ERROR",
             _("No DNS record found for %s")
             % "does-not-exist.example.com"),
            ("modoboa.admin", "ERROR",
             _("No DNS record found for %s")
             % "does-not-exist.example.com"),
            ("modoboa.admin", "WARNING",
             _("Invalid IP address format for %(domain)s; %(addr)s")
             % {"domain": "bad-response.example.com", "addr": "000.0.0.0"}),
            ("modoboa.admin", "WARNING",
             _("Invalid IP address format for %(domain)s; %(addr)s")
             % {"domain": "bad-response.example.com", "addr": "000.0.0.0"}),
        )

@override_settings(DNSBL_PROVIDERS=["zen.spamhaus.org"])
class DNSBLTestCase(ModoTestCase):
    """TestCase for DNSBL related features."""

    @classmethod
    def setUpTestData(cls):  # NOQA:N802
        """Create some data."""
        super().setUpTestData()
        cls.domain = factories.DomainFactory(name="modoboa.org")
        factories.DomainFactory(name="does-not-exist.example.com")
        cls.domain2 = factories.DomainFactory(
            name="test.localhost")  # Should not be checked
        cls.domain3 = factories.DomainFactory(
            name="modoboa.com", enable_dns_checks=False)
        models.DNSBLResult.objects.all().delete()

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        self.assertEqual(models.DNSBLResult.objects.count(), 0)
        management.call_command("modo", "check_mx")
        self.assertTrue(
            models.DNSBLResult.objects.filter(domain=self.domain).exists())
        self.assertFalse(
            models.DNSBLResult.objects.filter(domain=self.domain3).exists())
        response = self.client.get(
            reverse("admin:dnsbl_domain_detail", args=[self.domain.pk]))
        self.assertEqual(response.status_code, 200)
        self.assertFalse(self.domain.uses_a_reserved_tld)
        self.assertTrue(self.domain2.uses_a_reserved_tld)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_notifications(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check notifications."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "127.0.0.4"
        management.call_command(
            "modo", "check_mx", "--email", "user@example.test")
        self.assertEqual(len(mail.outbox), 2)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_notifications_wrong_dnsbl_response(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check notifications."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "127.255.255.254" #<--Spamhaus response when querying from an open resolver
        management.call_command(
            "modo", "check_mx", "--email", "user@example.test")
        self.assertEqual(len(mail.outbox), 1)

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command_no_dnsbl(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine without dnsbl."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"
        self.assertEqual(models.DNSBLResult.objects.count(), 0)
        management.call_command("modo", "check_mx", "--no-dnsbl")
        self.assertFalse(
            models.DNSBLResult.objects.filter(domain=self.domain).exists())
        response = self.client.get(
            reverse("admin:dnsbl_domain_detail", args=[self.domain.pk]))
        self.assertEqual(response.status_code, 200)

class DNSChecksTestCase(ModoTestCase):
    """A test case for DNS checks."""

    @classmethod
    def setUpTestData(cls):  # NOQA:N802
        """Create some data."""
        super(DNSChecksTestCase, cls).setUpTestData()
        cls.domain = factories.DomainFactory(name="dns-checks.com")

    @mock.patch("gevent.socket.gethostbyname")
    @mock.patch("socket.getaddrinfo")
    @mock.patch.object(dns.resolver.Resolver, "resolve")
    def test_management_command(
            self, mock_query, mock_getaddrinfo, mock_g_gethostbyname):
        """Check that command works fine."""
        mock_query.side_effect = utils.mock_dns_query_result
        mock_getaddrinfo.side_effect = utils.mock_ip_query_result
        mock_g_gethostbyname.return_value = "1.2.3.4"

        self.domain.enable_dkim = True
        self.domain.dkim_public_key = "XXXXX"
        self.domain.save(update_fields=["enable_dkim", "dkim_public_key"])

        management.call_command("modo", "check_mx", "--no-dnsbl", "--ttl=0")

        self.assertIsNot(self.domain.spf_record, None)
        self.assertIsNot(self.domain.dkim_record, None)
        self.assertIsNot(self.domain.dmarc_record, None)
        self.assertIsNot(self.domain.autoconfig_record, None)
        self.assertIsNot(self.domain.autodiscover_record, None)
<|endoftext|>


<commit_before>from files import find
def init(args):

	path = args["root"] if "root" in args else "."
	(rootdir, file) = find("strings.xml", path)

	global root
	root = rootdir

	global resourceFiles
	resourceFiles = {}

def close():
	for storyboard, file in resourceFiles.iteritems():
		file.write("\n</resources>")
		file.close()

def writeComment(row, storyboard):
	if storyboard != None:
		file = resourceFiles[storyboard.lower()]
		if file != None:
			file.write("\n\t\n")

def writeText(row, storyboard):
	if storyboard != None:
		file = resourceFiles[storyboard.lower()]
		if file != None:
			file.write("\t<string name=\""+row.key.strip()+"\">"+row.text.encode('utf-8').replace("
","\n")+"</string>\n")

def writeStoryboard(row, storyboard):
	if root != None:
		file = open(root+"/strings-gen-"+storyboard.lower()+".xml", "wb+")
		resourceFiles[storyboard.lower()] = file
		file.write("<resources>\n\t\n")<commit_msg>Escape `\n` and `'` characters on the generated string (android only)<commit_after>from files import find
def init(args):

	path = args["root"] if "root" in args else "."
	(rootdir, file) = find("strings.xml", path)

	global root
	root = rootdir

	global resourceFiles
	resourceFiles = {}

def close():
	for storyboard, file in resourceFiles.iteritems():
		file.write("\n</resources>")
		file.close()

def writeComment(row, storyboard):
	if storyboard != None:
		file = resourceFiles[storyboard.lower()]
		if file != None:
			file.write("\n\t\n")

def writeText(row, storyboard):
	if storyboard != None:
		file = resourceFiles[storyboard.lower()]
		if file != None:
			content = row.text.encode('utf-8').replace("
","\\n").replace("'", "\\'")
			file.write("\t<string name=\""+row.key.strip()+"\">"+content+"</string>\n")

def writeStoryboard(row, storyboard):
	if root != None:
		file = open(root+"/strings-gen-"+storyboard.lower()+".xml", "wb+")
		resourceFiles[storyboard.lower()] = file
		file.write("<resources>\n\t\n")
<|endoftext|>


<commit_before>import time
from threading import Thread

from cassandra import ConsistencyLevel
from ccmlib.node import NodetoolError

from dtest import Tester
from tools import insert_c1c2, known_failure, query_c1c2, since

class TestRebuild(Tester):

    def __init__(self, *args, **kwargs):
        kwargs['cluster_options'] = {'start_rpc': 'true'}
        # Ignore these log patterns:
        self.ignore_log_patterns = [
            # This one occurs when trying to send the migration to a
            # node that hasn't started yet, and when it does, it gets
            # replayed and everything is fine.
            r'Can\'t send migration request: node.*is down',
            # ignore streaming error during bootstrap
            r'Exception encountered during startup',
            r'Streaming error occurred'
        ]
        Tester.__init__(self, *args, **kwargs)

    @known_failure(failure_source='test',
                   jira_url='https://issues.apache.org/jira/browse/CASSANDRA-11687',
                   flaky=True)
    def simple_rebuild_test(self):
        """
        @jira_ticket CASSANDRA-9119

        Test rebuild from other dc works as expected.
        """

        keys = 1000

        cluster = self.cluster
        cluster.set_configuration_options(values={'endpoint_snitch': 'org.apache.cassandra.locator.PropertyFileSnitch'})
        node1 = cluster.create_node('node1', False,
                                    ('127.0.0.1', 9160),
                                    ('127.0.0.1', 7000),
                                    '7100', '2000', None,
                                    binary_interface=('127.0.0.1', 9042))
        cluster.add(node1, True, data_center='dc1')

        # start node in dc1
        node1.start(wait_for_binary_proto=True)

        # populate data in dc1
        session = self.patient_exclusive_cql_connection(node1)
        self.create_ks(session, 'ks', {'dc1': 1})
        self.create_cf(session, 'cf', columns={'c1': 'text', 'c2': 'text'})
        insert_c1c2(session, n=keys, consistency=ConsistencyLevel.ALL)

        # check data
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ALL)
        session.shutdown()

        # Bootstraping a new node in dc2 with auto_bootstrap: false
        node2 = cluster.create_node('node2', False,
                                    ('127.0.0.2', 9160),
                                    ('127.0.0.2', 7000),
                                    '7200', '2001', None,
                                    binary_interface=('127.0.0.2', 9042))
        cluster.add(node2, False, data_center='dc2')
        node2.start(wait_other_notice=True, wait_for_binary_proto=True)

        # wait for snitch to reload
        time.sleep(60)
        # alter keyspace to replicate to dc2
        session = self.patient_exclusive_cql_connection(node2)
        session.execute("ALTER KEYSPACE ks WITH REPLICATION = {'class':'NetworkTopologyStrategy', 'dc1':1, 'dc2':1};")
        session.execute('USE ks')

        self.rebuild_errors = 0
        # make it possible for thread to raise errors
        self.thread_exc_info = None

        # rebuild dc2 from dc1
        def rebuild():
            try:
                node2.nodetool('rebuild dc1')
            except NodetoolError as e:
                if 'Node is still rebuilding' in e.message:
                    self.rebuild_errors += 1
                else:
                    raise e

        def call_and_register_exc_info_on_exception(func):
            """
            Closes over self to catch any exceptions raised by func and
            register them at self.thread_exc_info
            Based on http://stackoverflow.com/a/1854263
            """
            try:
                func()
            except Exception:
                import sys
                self.thread_exc_info = sys.exc_info()

        cmd1 = Thread(target=call_and_register_exc_info_on_exception(rebuild))
        cmd1.start()

        # concurrent rebuild should not be allowed (CASSANDRA-9119)
        # (following sleep is needed to avoid conflict in 'nodetool()' method setting up env.)
        time.sleep(.1)
        # we don't need to manually raise exeptions here -- already handled
        rebuild()

        cmd1.join()

        # manually raise exception from cmd1 thread
        # see http://stackoverflow.com/a/1854263
        if self.thread_exc_info is not None:
            raise self.thread_exc_info[1], None, self.thread_exc_info[2]

        # exactly 1 of the two nodetool calls should fail
        # usually it will be the one in the main thread,
        # but occasionally it wins the race with the one in the secondary thread,
        # so we check that one succeeded and the other failed
        self.assertEqual(self.rebuild_errors, 1,
                         msg='rebuild errors should be 1, but found {}. Concurrent rebuild should not be allowed, but one rebuild command should have succeeded.'.format(self.rebuild_errors))

        # check data
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ALL)

    @since('3.6')
    def rebuild_ranges_test(self):
        """
        @jira_ticket CASSANDRA-10406
        """
        keys = 1000

        cluster = self.cluster
        tokens = cluster.balanced_tokens_across_dcs(['dc1', 'dc2'])
        cluster.set_configuration_options(values={'endpoint_snitch': 'org.apache.cassandra.locator.PropertyFileSnitch'})
        cluster.set_configuration_options(values={'num_tokens': 1})
        node1 = cluster.create_node('node1', False,
                                    ('127.0.0.1', 9160),
                                    ('127.0.0.1', 7000),
                                    '7100', '2000', tokens[0],
                                    binary_interface=('127.0.0.1', 9042))
        node1.set_configuration_options(values={'initial_token': tokens[0]})
        cluster.add(node1, True, data_center='dc1')
        node1 = cluster.nodelist()[0]

        # start node in dc1
        node1.start(wait_for_binary_proto=True)

        # populate data in dc1
        session = self.patient_exclusive_cql_connection(node1)
        # ks1 will be rebuilt in node2
        self.create_ks(session, 'ks1', {'dc1': 1})
        self.create_cf(session, 'cf', columns={'c1': 'text', 'c2': 'text'})
        insert_c1c2(session, n=keys, consistency=ConsistencyLevel.ALL)
        # ks2 will not be rebuilt in node2
        self.create_ks(session, 'ks2', {'dc1': 1})
        self.create_cf(session, 'cf', columns={'c1': 'text', 'c2': 'text'})
        insert_c1c2(session, n=keys, consistency=ConsistencyLevel.ALL)
        session.shutdown()

        # Bootstraping a new node in dc2 with auto_bootstrap: false
        node2 = cluster.create_node('node2', False,
                                    ('127.0.0.2', 9160),
                                    ('127.0.0.2', 7000),
                                    '7200', '2001', tokens[1],
                                    binary_interface=('127.0.0.2', 9042))
        node2.set_configuration_options(values={'initial_token': tokens[1]})
        cluster.add(node2, False, data_center='dc2')
        node2.start(wait_other_notice=True, wait_for_binary_proto=True)

        # wait for snitch to reload
        time.sleep(60)
        # alter keyspace to replicate to dc2
        session = self.patient_exclusive_cql_connection(node2)
        session.execute("ALTER KEYSPACE ks1 WITH REPLICATION = {'class':'NetworkTopologyStrategy', 'dc1':1, 'dc2':1};")
        session.execute("ALTER KEYSPACE ks2 WITH REPLICATION = {'class':'NetworkTopologyStrategy', 'dc1':1, 'dc2':1};")
        session.execute('USE ks1')

        # rebuild only ks1 with range that is node1's replica
        node2.nodetool('rebuild -ks ks1 -ts (%s,%s] dc1' % (tokens[1], str(pow(2, 63) - 1)))

        # check data is sent by stopping node1
        node1.stop()
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ONE)
        # ks2 should not be streamed
        session.execute('USE ks2')
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ONE, tolerate_missing=True, must_be_missing=True)
<commit_msg>Fix typo<commit_after>import time
from threading import Thread

from cassandra import ConsistencyLevel
from ccmlib.node import NodetoolError

from dtest import Tester
from tools import insert_c1c2, known_failure, query_c1c2, since

class TestRebuild(Tester):

    def __init__(self, *args, **kwargs):
        kwargs['cluster_options'] = {'start_rpc': 'true'}
        # Ignore these log patterns:
        self.ignore_log_patterns = [
            # This one occurs when trying to send the migration to a
            # node that hasn't started yet, and when it does, it gets
            # replayed and everything is fine.
            r'Can\'t send migration request: node.*is down',
            # ignore streaming error during bootstrap
            r'Exception encountered during startup',
            r'Streaming error occurred'
        ]
        Tester.__init__(self, *args, **kwargs)

    @known_failure(failure_source='test',
                   jira_url='https://issues.apache.org/jira/browse/CASSANDRA-11687',
                   flaky=True)
    def simple_rebuild_test(self):
        """
        @jira_ticket CASSANDRA-9119

        Test rebuild from other dc works as expected.
        """

        keys = 1000

        cluster = self.cluster
        cluster.set_configuration_options(values={'endpoint_snitch': 'org.apache.cassandra.locator.PropertyFileSnitch'})
        node1 = cluster.create_node('node1', False,
                                    ('127.0.0.1', 9160),
                                    ('127.0.0.1', 7000),
                                    '7100', '2000', None,
                                    binary_interface=('127.0.0.1', 9042))
        cluster.add(node1, True, data_center='dc1')

        # start node in dc1
        node1.start(wait_for_binary_proto=True)

        # populate data in dc1
        session = self.patient_exclusive_cql_connection(node1)
        self.create_ks(session, 'ks', {'dc1': 1})
        self.create_cf(session, 'cf', columns={'c1': 'text', 'c2': 'text'})
        insert_c1c2(session, n=keys, consistency=ConsistencyLevel.ALL)

        # check data
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ALL)
        session.shutdown()

        # Bootstrapping a new node in dc2 with auto_bootstrap: false
        node2 = cluster.create_node('node2', False,
                                    ('127.0.0.2', 9160),
                                    ('127.0.0.2', 7000),
                                    '7200', '2001', None,
                                    binary_interface=('127.0.0.2', 9042))
        cluster.add(node2, False, data_center='dc2')
        node2.start(wait_other_notice=True, wait_for_binary_proto=True)

        # wait for snitch to reload
        time.sleep(60)
        # alter keyspace to replicate to dc2
        session = self.patient_exclusive_cql_connection(node2)
        session.execute("ALTER KEYSPACE ks WITH REPLICATION = {'class':'NetworkTopologyStrategy', 'dc1':1, 'dc2':1};")
        session.execute('USE ks')

        self.rebuild_errors = 0
        # make it possible for thread to raise errors
        self.thread_exc_info = None

        # rebuild dc2 from dc1
        def rebuild():
            try:
                node2.nodetool('rebuild dc1')
            except NodetoolError as e:
                if 'Node is still rebuilding' in e.message:
                    self.rebuild_errors += 1
                else:
                    raise e

        def call_and_register_exc_info_on_exception(func):
            """
            Closes over self to catch any exceptions raised by func and
            register them at self.thread_exc_info
            Based on http://stackoverflow.com/a/1854263
            """
            try:
                func()
            except Exception:
                import sys
                self.thread_exc_info = sys.exc_info()

        cmd1 = Thread(target=call_and_register_exc_info_on_exception(rebuild))
        cmd1.start()

        # concurrent rebuild should not be allowed (CASSANDRA-9119)
        # (following sleep is needed to avoid conflict in 'nodetool()' method setting up env.)
        time.sleep(.1)
        # we don't need to manually raise exeptions here -- already handled
        rebuild()

        cmd1.join()

        # manually raise exception from cmd1 thread
        # see http://stackoverflow.com/a/1854263
        if self.thread_exc_info is not None:
            raise self.thread_exc_info[1], None, self.thread_exc_info[2]

        # exactly 1 of the two nodetool calls should fail
        # usually it will be the one in the main thread,
        # but occasionally it wins the race with the one in the secondary thread,
        # so we check that one succeeded and the other failed
        self.assertEqual(self.rebuild_errors, 1,
                         msg='rebuild errors should be 1, but found {}. Concurrent rebuild should not be allowed, but one rebuild command should have succeeded.'.format(self.rebuild_errors))

        # check data
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ALL)

    @since('3.6')
    def rebuild_ranges_test(self):
        """
        @jira_ticket CASSANDRA-10406
        """
        keys = 1000

        cluster = self.cluster
        tokens = cluster.balanced_tokens_across_dcs(['dc1', 'dc2'])
        cluster.set_configuration_options(values={'endpoint_snitch': 'org.apache.cassandra.locator.PropertyFileSnitch'})
        cluster.set_configuration_options(values={'num_tokens': 1})
        node1 = cluster.create_node('node1', False,
                                    ('127.0.0.1', 9160),
                                    ('127.0.0.1', 7000),
                                    '7100', '2000', tokens[0],
                                    binary_interface=('127.0.0.1', 9042))
        node1.set_configuration_options(values={'initial_token': tokens[0]})
        cluster.add(node1, True, data_center='dc1')
        node1 = cluster.nodelist()[0]

        # start node in dc1
        node1.start(wait_for_binary_proto=True)

        # populate data in dc1
        session = self.patient_exclusive_cql_connection(node1)
        # ks1 will be rebuilt in node2
        self.create_ks(session, 'ks1', {'dc1': 1})
        self.create_cf(session, 'cf', columns={'c1': 'text', 'c2': 'text'})
        insert_c1c2(session, n=keys, consistency=ConsistencyLevel.ALL)
        # ks2 will not be rebuilt in node2
        self.create_ks(session, 'ks2', {'dc1': 1})
        self.create_cf(session, 'cf', columns={'c1': 'text', 'c2': 'text'})
        insert_c1c2(session, n=keys, consistency=ConsistencyLevel.ALL)
        session.shutdown()

        # Bootstraping a new node in dc2 with auto_bootstrap: false
        node2 = cluster.create_node('node2', False,
                                    ('127.0.0.2', 9160),
                                    ('127.0.0.2', 7000),
                                    '7200', '2001', tokens[1],
                                    binary_interface=('127.0.0.2', 9042))
        node2.set_configuration_options(values={'initial_token': tokens[1]})
        cluster.add(node2, False, data_center='dc2')
        node2.start(wait_other_notice=True, wait_for_binary_proto=True)

        # wait for snitch to reload
        time.sleep(60)
        # alter keyspace to replicate to dc2
        session = self.patient_exclusive_cql_connection(node2)
        session.execute("ALTER KEYSPACE ks1 WITH REPLICATION = {'class':'NetworkTopologyStrategy', 'dc1':1, 'dc2':1};")
        session.execute("ALTER KEYSPACE ks2 WITH REPLICATION = {'class':'NetworkTopologyStrategy', 'dc1':1, 'dc2':1};")
        session.execute('USE ks1')

        # rebuild only ks1 with range that is node1's replica
        node2.nodetool('rebuild -ks ks1 -ts (%s,%s] dc1' % (tokens[1], str(pow(2, 63) - 1)))

        # check data is sent by stopping node1
        node1.stop()
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ONE)
        # ks2 should not be streamed
        session.execute('USE ks2')
        for i in xrange(0, keys):
            query_c1c2(session, i, ConsistencyLevel.ONE, tolerate_missing=True, must_be_missing=True)
<|endoftext|>


<commit_before>import logging
import logging.handlers
import graypy
import os
import sys

from .config import config

class StreamToLogger(object):
    """
    Fake file-like stream object that redirects writes to a logger instance.
    """

    def __init__(self, logger, log_level=logging.INFO):
        self.logger = logger
        self.log_level = log_level
        self.linebuf = ''

    def write(self, buf):
        for line in buf.rstrip().splitlines():
            self.logger.log(self.log_level, line.rstrip())

def setup_logging(logdir=None, scrnlog=True, txtlog=True, loglevel=logging.DEBUG):
    logdir = os.path.abspath(logdir)

    if not os.path.exists(logdir):
        os.mkdir(logdir)

    log = logging.getLogger('')
    log.setLevel(loglevel)

    log_formatter = logging.Formatter("%(asctime)s - %(levelname)-7s :: %(name)-6s :: %(message)s")

    if hasattr(config, 'GRAYLOG'):
        graylog_handler = graypy.GELFHandler(config.GRAYLOG)
        graylog_handler.setFormatter(log_formatter)
        log.addHandler(graylog_handler)

    if txtlog:
        txt_handler = logging.handlers.RotatingFileHandler(os.path.join(logdir, "vmmaster.log"), backupCount=5)
        txt_handler.setFormatter(log_formatter)
        log.addHandler(txt_handler)

        log.info("Logger initialised.")

    if scrnlog:
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(log_formatter)
        log.addHandler(console_handler)

    stdout_logger = logging.getLogger('STDOUT')
    slout = StreamToLogger(stdout_logger, logging.INFO)
    sys.stdout = slout

    stderr_logger = logging.getLogger('STDERR')
    slerr = StreamToLogger(stderr_logger, logging.ERROR)
    sys.stderr = slerr

    return log

log = logging.getLogger('LOG')<commit_msg>fix log rotation size<commit_after>import logging
import logging.handlers
import graypy
import os
import sys

from .config import config

class StreamToLogger(object):
    """
    Fake file-like stream object that redirects writes to a logger instance.
    """

    def __init__(self, logger, log_level=logging.INFO):
        self.logger = logger
        self.log_level = log_level
        self.linebuf = ''

    def write(self, buf):
        for line in buf.rstrip().splitlines():
            self.logger.log(self.log_level, line.rstrip())

def setup_logging(logdir=None, scrnlog=True, txtlog=True, loglevel=logging.DEBUG):
    logdir = os.path.abspath(logdir)

    if not os.path.exists(logdir):
        os.mkdir(logdir)

    log = logging.getLogger('')
    log.setLevel(loglevel)

    log_formatter = logging.Formatter("%(asctime)s - %(levelname)-7s :: %(name)-6s :: %(message)s")

    if hasattr(config, 'GRAYLOG'):
        graylog_handler = graypy.GELFHandler(config.GRAYLOG)
        graylog_handler.setFormatter(log_formatter)
        log.addHandler(graylog_handler)

    if txtlog:
        txt_handler = logging.handlers.RotatingFileHandler(
            os.path.join(logdir, "vmmaster.log"), maxBytes=5120, backupCount=5
        )
        txt_handler.setFormatter(log_formatter)
        log.addHandler(txt_handler)

        log.info("Logger initialised.")

    if scrnlog:
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(log_formatter)
        log.addHandler(console_handler)

    stdout_logger = logging.getLogger('STDOUT')
    slout = StreamToLogger(stdout_logger, logging.INFO)
    sys.stdout = slout

    stderr_logger = logging.getLogger('STDERR')
    slerr = StreamToLogger(stderr_logger, logging.ERROR)
    sys.stderr = slerr

    return log

log = logging.getLogger('LOG')<|endoftext|>


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

# Copyright(C) 2010-2012 Romain Bignon, Florent Fourcot
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

from weboob.tools.browser2.page import HTMLPage, method, ListElement, ItemElement
from weboob.tools.browser2.filters import CleanText, Env, Regexp, Field
from weboob.capabilities.gauge import GaugeMeasure, GaugeSensor
from weboob.capabilities.base import NotAvailable

__all__ = ['StartPage']

class StartPage(HTMLPage):

    @method
    class get_sensors_list(ListElement):
        item_xpath = '//p[@align="center"]'

        class item(ItemElement):
            klass = GaugeSensor

            obj_name = Regexp(CleanText('.'), '(.*?) {0,}: .*', "\\1")
            obj_id = CleanText(Regexp(Field('name'), '(.*)', "dd-\\1"), " .():")
            obj_unit = Env('unit')
            obj_lastvalue = Env('lastvalue')
            obj_gaugeid = u"wetter"
            obj_forecast = NotAvailable

            def split_unit(self, text):
                if u"Temperatur" in text:
                    value = text.split(': ')[1].split(u'°')[0]
                    unit = u'°C'
                else:
                    value = text.split(':')[-1].split()[0]
                    unit = text.split(':')[-1].split(None, 1)[1]
                return value, unit

            def parse(self, el):
                text = CleanText(el)(self)

                level, self.env['unit'] = self.split_unit(text)
                lastvalue = GaugeMeasure()
                lastvalue.level = float(level)
                lastvalue.alarm = NotAvailable
                self.env['lastvalue'] = lastvalue
<commit_msg>Simplify parse function<commit_after># -*- coding: utf-8 -*-

# Copyright(C) 2010-2012 Romain Bignon, Florent Fourcot
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

from weboob.tools.browser2.page import HTMLPage, method, ListElement, ItemElement
from weboob.tools.browser2.filters import CleanText, Env, Regexp, Field
from weboob.capabilities.gauge import GaugeMeasure, GaugeSensor
from weboob.capabilities.base import NotAvailable

__all__ = ['StartPage']

class StartPage(HTMLPage):

    @method
    class get_sensors_list(ListElement):
        item_xpath = '//p[@align="center"]'

        class item(ItemElement):
            klass = GaugeSensor

            obj_name = Regexp(CleanText('.'), '(.*?) {0,}: .*', "\\1")
            obj_id = CleanText(Regexp(Field('name'), '(.*)', "dd-\\1"), " .():")
            obj_gaugeid = u"wetter"
            obj_forecast = NotAvailable

            def split_unit(self, text):
                if u"Temperatur" in text:
                    value = text.split(': ')[1].split(u'°')[0]
                    unit = u'°C'
                else:
                    value = text.split(':')[-1].split()[0]
                    unit = text.split(':')[-1].split(None, 1)[1]
                return [value, unit]

            def obj_lastvalue(self):
                lastvalue = GaugeMeasure()
                lastvalue.level = float(Env('parse')(self)[0])
                lastvalue.alarm = NotAvailable
                return lastvalue

            def obj_unit(self):
                return Env('parse')(self)[1]

            def parse(self, el):
                self.env['parse'] = self.split_unit(CleanText(el)(self))
<|endoftext|>


<commit_before>"""
Django settings for voting_site project.

For more information on this file, see
https://docs.djangoproject.com/en/1.7/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/1.7/ref/settings/
"""

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
import os
BASE_DIR = os.path.dirname(os.path.dirname(__file__))

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/1.7/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = '$s!h$4hp-f@l64kr2v(x&=&)kwn@5ly%wcry^97&_mge!pt-_5'

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = True

TEMPLATE_DEBUG = True

ALLOWED_HOSTS = []

# Application definition

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'crispy_forms',
    'voting',
)

CRISPY_TEMPLATE_PACK= 'bootstrap3'

MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
    'django.middleware.locale.LocaleMiddleware',
)

ROOT_URLCONF = 'voting_site.urls'

WSGI_APPLICATION = 'voting_site.wsgi.application'

# Database
# https://docs.djangoproject.com/en/1.7/ref/settings/#databases

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

# Internationalization
# https://docs.djangoproject.com/en/1.7/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'Asia/Taipei'

USE_I18N = True

USE_L10N = True

USE_TZ = False

LOCALE_PATHS = (
    os.path.join(BASE_DIR, 'locale'),
)

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/1.7/howto/static-files/

STATIC_URL = '/static/'

# login/logout
LOGIN_URL = '/admin/login'
LOGOUT_URL = '/admin/logout'
<commit_msg>Add STATIC_ROOT setting<commit_after>"""
Django settings for voting_site project.

For more information on this file, see
https://docs.djangoproject.com/en/1.7/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/1.7/ref/settings/
"""

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
import os
BASE_DIR = os.path.dirname(os.path.dirname(__file__))

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/1.7/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = '$s!h$4hp-f@l64kr2v(x&=&)kwn@5ly%wcry^97&_mge!pt-_5'

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = True

TEMPLATE_DEBUG = True

ALLOWED_HOSTS = []

# Application definition

INSTALLED_APPS = (
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'crispy_forms',
    'voting',
)

CRISPY_TEMPLATE_PACK= 'bootstrap3'

MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
    'django.middleware.locale.LocaleMiddleware',
)

ROOT_URLCONF = 'voting_site.urls'

WSGI_APPLICATION = 'voting_site.wsgi.application'

# Database
# https://docs.djangoproject.com/en/1.7/ref/settings/#databases

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

# Internationalization
# https://docs.djangoproject.com/en/1.7/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'Asia/Taipei'

USE_I18N = True

USE_L10N = True

USE_TZ = False

LOCALE_PATHS = (
    os.path.join(BASE_DIR, 'locale'),
)

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/1.7/howto/static-files/

STATIC_URL = '/static/'
STATIC_ROOT = os.path.join(BASE_DIR, 'public')

# login/logout
LOGIN_URL = '/admin/login'
LOGOUT_URL = '/admin/logout'
<|endoftext|>


<commit_before>#!/usr/bin/python
# This fixes some of the image URLs in records.pickle.
import record
import re
import cPickle
rs = record.AllRecords()

for idx, r in enumerate(rs):
  url = r.photo_url

  # Some images have thumbnails but are missing the full photo URL.
  # Convert
  #    http://webbie1.sfpl.org/multimedia/thumbnails/aaa9774_x.jpg
  # -> http://webbie1.sfpl.org/multimedia/sfphotos/aaa-9774.jpg
  if not url:
    url = r.thumbnail_url
    if 'aaf' in url:
      # There are only two of these:
      #   http://sflib1.sfpl.org:82/record=b1026391~S0
      #   http://sflib1.sfpl.org:82/record=b1036043~S0
      # They're both links to larger collections of images and don't really fit
      # the mold of the other pages. So we remove them.
      del rs[idx]
      continue
    else:
      # Two of these... I think they're just omissions.
      url = re.sub(r'(.*)/thumbnails/(...)(\d+)_x\.jpg', r'\1/sfphotos/\2-\3.jpg', url)

  # Remove trailing spaces from image URLs.
  # Maybe ~4 of these.
  if url[-1] == ' ':
    url = url[0:-1]

  # Change 'foojpg' -> 'foo.jpg'
  # Just a typo. There are ~4 of these, too.
  if url[-3:] == 'jpg' and url[-4:] != '.jpg':
    url = url[:-3] + '.jpg'

print "Re-pickling"
output_file = "records.pickle"
f = file(output_file, "w")
p = cPickle.Pickler(f, 2)
for r in rs:
  count += 1
  p.dump(r)
  if count % 100 == 0:
    print "Pickled %d records" % count
<commit_msg>fix bug<commit_after>#!/usr/bin/python
# This fixes some of the image URLs in records.pickle.
import record
import re
import cPickle
rs = record.AllRecords()

for idx, r in enumerate(rs):
  url = r.photo_url

  # Some images have thumbnails but are missing the full photo URL.
  # Convert
  #    http://webbie1.sfpl.org/multimedia/thumbnails/aaa9774_x.jpg
  # -> http://webbie1.sfpl.org/multimedia/sfphotos/aaa-9774.jpg
  if not url:
    url = r.thumbnail_url
    if 'aaf' in url:
      # There are only two of these:
      #   http://sflib1.sfpl.org:82/record=b1026391~S0
      #   http://sflib1.sfpl.org:82/record=b1036043~S0
      # They're both links to larger collections of images and don't really fit
      # the mold of the other pages. So we remove them.
      del rs[idx]
      continue
    else:
      # Two of these... I think they're just omissions.
      url = re.sub(r'(.*)/thumbnails/(...)(\d+)_x\.jpg', r'\1/sfphotos/\2-\3.jpg', url)

  # Remove trailing spaces from image URLs.
  # Maybe ~4 of these.
  if url[-1] == ' ':
    url = url[0:-1]

  # Change 'foojpg' -> 'foo.jpg'
  # Just a typo. There are ~4 of these, too.
  if url[-3:] == 'jpg' and url[-4:] != '.jpg':
    url = url[:-3] + '.jpg'

print "Re-pickling"
output_file = "records.pickle"
f = file(output_file, "w")
p = cPickle.Pickler(f, 2)
for i, r in enumerate(rs):
  p.dump(r)
  if i % 100 == 0:
    print "Pickled %d records" % i
<|endoftext|>


<commit_before>"""

.. module:: bigwig
    :synopsis: BIGWIG file related stuffs

..  moduleauthor:: Ken Sugino <ken.sugino@gmail.com>

"""

# system
import os
import subprocess
import time
import multiprocessing
from operator import iadd
import logging
logging.basicConfig(level=logging.DEBUG)
LOG = logging.getLogger(__name__)
import shutil

# 3rd party
import pandas as PD
import numpy as N

# for reading BIGWIG file
#ngslib is faster than bx but fails when parallerized
#import bx
#from bx.bbi.bigwig_file import BigWigFile
#from ngslib import wWigIO
#bx-python is not compatible with Python3 => use modified version

from jgem import utils as UT
from jgem import bedtools as BT
import jgem.cy.bw  as cybw #import array2wiggle_chr # Cython version
from jgem.cy.bw import array2wiggle_chr
from jgem.bxbbi.bigwig_file import BigWigFile

MAXSIZE = int(300e6)  # 300Mbp bigger than chr1,chrX

# BAM to genomecov BIGWIG ##################################################
def cnt_bam(fpath):
    """
    Uses samtools to count reads in BAM file.

    Args:
        fpath (str): path to BAM file

    Returns:
        int. The number of aligned reads.

    Raises:
        RuntimeError
    """
    cache = fpath+'.flagstat'
    bai = fpath+'.bai'
    if not os.path.exists(fpath):
        raise RuntimeError('file {0} does not exist'.format(fpath))
    if not os.path.exists(bai):
        cmd = ['samtools', 'index', fpath]
        subprocess.call(cmd)
    if os.path.exists(cache) and os.path.getmtime(cache)>os.path.getmtime(fpath):
        out = open(cache,'r').read()
    else:
        cmd = ['samtools', 'flagstat', fpath]
        out = subprocess.check_output(cmd)
        open(cache,'wb').write(out)
    firstline = out.split('\n')[0].split()
    return int(firstline[0])+int(firstline[2])

def wig2bw(wigpath, chromsizes, bwpath):
    pass

def bam2bw(fpath, chromsizes, bpath, aligned=None):
    """
    Generate normalized coverage from BAM

    Args:
        fpath (str): path to BAM
        chromsizes (str): path to chromsizes file
        bpath (str): path to BIGWIG
        aligned (int): number of aligned reads, if None uses samtools to find it from BAM

    Requires Bedtools (genomeCoverageBed) and Kent Tool (wigToBigWig)

    """
    # countreads
    if aligned is None:
        aligned =   cnt_bam(fpath)
    scale = 1000000./float(aligned)
    # convert_to_wig
    tpath = bpath +'.wig'
    UT.makedirs(os.path.dirname(tpath))
    tfobj = open(tpath,'wb')
    cmd1 = ['genomeCoverageBed', '-split', '-bg', '-ibam', fpath, '-g', chromsizes, '-scale', str(scale)]
    p1 = subprocess.Popen(cmd1, stdout=tfobj)
    p1.wait()
    tfobj.close()

    # convet_wig_to_bigwig
    cmd2 = ['wigToBigWig', tpath, chromsizes, bpath]
    p2 = subprocess.call(cmd2)

    # remove_temporary_file
    os.remove(tpath)

# bw2bed based on #########################################################
# https://github.com/CGATOxford/cgat/blob/master/scripts/wig2bed.py
# [TODO] report bug to AndreasHeger line 80 "elif => if"

def block_iter(infile, chrom, chunk=int(10e6)):
    "BigWig file iterator"
    with open(infile, mode='rb') as fobj:
        bw = BigWigFile(fobj)
        for x in range(0,MAXSIZE,chunk): # 10Mbp chunk
            iterator = bw.get(chrom, x, x+chunk)
            if iterator is None:
                raise StopIteration
            for v in iterator:
                yield v

# def block_iter_ngs(infile, chrom, chunk=int(10e6)):
#     # ngslib is faster than bx but fails when parallerized
#     wWigIO.open(infile)
#     for x in range(0,MAXSIZE,chunk): # 10Mbp chunk
#         iterator = wWigIO.getIntervals(infile, chrom, x, x+chunk)
#         if not iterator:
#             raise StopIteration
#         for v in iterator:
#             yield v
#     wWigIO.close(infile)

def apply_threshold(infile, threshold, chroms):
    """Only returns intervals exceeding the threshold

    Args:
        infile: path to bigwig
        threshold: positive float
        chroms: list of chromosome names

    Yields:
        intervals (chr,st,ed)
    """
    for chrom in chroms:
        last_start, last_end = -1, 0
        #for start, end, value in block_iter_ngs(infile, chrom):
        for start, end, value in block_iter(infile, chrom):
            # start=>end has value
            d = start - last_end
            # d: distance from last_end
            # d> 0 => last_end==>start value is 0 assume threshold>=0
            if (d > 0 or value <= threshold):
                if last_start >= 0: # if there's un-yielded interval then yield
                    yield chrom, last_start, last_end
                last_start = -1 # reset (no hit)
            #elif last_start < 0 and value > threshold:
            # this part is a bug from the original code
            # original code will skip the next interval
            if last_start < 0 and value > threshold:
                # last_start <0 => no current interval
                # and value above threshod ==> start new interval
                last_start = start
            last_end = end
        if last_start >= 0:
            yield chrom, last_start, end

def bw2bed(bwfile, bedfile, chroms, th, compress=True):
    """Transform BigWig genomeCov to binary BED by thresholding.
    Makes result file (bwfile[:-3]+'.binary%g.bed'.format(th))

    Args:
        bwfile: path to BigWig file
        chroms: list of chromosome names
        th: coverage threshold

    Returns:
        path to generated BED file
    """
    bedbase = bedfile[:-3] if bedfile[-3:]=='.gz' else bedfile
    #bedfile = '{0}.binary{1:g}.bed'.format(bwfile[:-3], th)
    if UT.notstale(bwfile, bedbase+'.gz'):
        return bedbase+'.gz'
    # make sure bwfile exists
    if not ( os.path.exists(bwfile) ):
        raise RuntimeError('BigWig file {0} does not exist.'.format(bwfile))
    processor = apply_threshold(bwfile,th, chroms)
    UT.makedirs(os.path.dirname(bedfile))
    out = open(bedbase,'w')
    out.write(''.join(['%s\t%i\t%i\n' % x for x in processor]))
    #out.write('\n') #<= this introduces space inbetween chroms in mp ode
    # which terminates bedtools at chr1
    out.close()
    if compress:
        return UT.compress(bedbase)
    return bedbase

def bw2bed_mp(bwfile, bedfile, chroms, th, np=4):
    """ multi CPU version of bw2bed """

    args = []
    files = []
    for chrom in chroms:
        bedchromfile = bedfile+'.{0}.bed.gz'.format(chrom)
        files.append(bedchromfile)
        args.append((bwfile,bedchromfile,[chrom],th,False))

    rslts = UT.process_mp(bw2bed, args, np=np, doreduce=False)

    # concatenate gz files
    bedbase = bedfile[:-3] if bedfile[-3:]=='.gz' else bedfile
    with open(bedbase, 'wb') as dst:
        for f in rslts:
            with open(f, 'rb') as src:
                shutil.copyfileobj(src, dst)
    # !!! bedtool gzip problem againg !!!
    # bedtools only process first one if just concatenate gzipped files
    # => concatenate unzipped and gzip whole thing at the end
    bedfile = UT.compress(bedbase)

    # clean up temp files
    for f in files:
        os.unlink(f)

    return bedfile

### covbp, totbp, avgcov from bigwig array #################################

def get_totbp_covbp_bw(bwfile, genome, chroms=None):
    """ Calculate total bp, covered bp, mean coverage, covered %.

    Args:
        bwfile: bigwig file
        genome: UCSC genome name
        chroms (list): of chromosomes

    Returns:
        Pandas dataframe

    """
    chromdf = UT.chromdf(genome).set_index('chr')['size']
    def one(chrom):
        csize = chromdf.ix[chrom]
        a = get_bigwig_as_array(bwfile, chrom, 0, csize)
        totbp = N.sum(a)
        covbp = N.sum(a>0)
        acov = float(totbp)/covbp
        covp = float(covbp)/csize
        return {'totbp':totbp,'covbp':covbp,'acov':acov,'cov%':covp}
    if chroms is None:
        chroms = UT.chroms(genome)
    df = PD.DataFrame({x: one(x) for x in chroms})
    return df

### Merge BigWigs ##########################################################

def get_bigwig_as_array(bwfile, chrom, st, ed):
    """Get BIGWIG coverage values as array of size (ed-st). Array start corresponds to st.
    0-based

    Args:
        bwfile: path to BIGWIG
        chrom (str): chromosome name
        st (int): start position
        ed (int): end position

    Returns:
        Numpy array of size (ed-st)
    """
    # with open(bwfile, mode='rb') as fobj:
    #     bw = BigWigFile(fobj)
    #     it = bw.get(chrom,st,ed)
    #     a = N.zeros(ed-st)
    #     for s,e,v in it:
    #         a[s-st:e-st] += v
    # return a
    with open(bwfile, mode='rb') as fobj:
        bw = BigWigFile(fobj)
        a = bw.get_as_array(chrom,st,ed)
        a[N.isnan(a)]=0.
    return a

def merge_bigwigs_chr(bwfiles, chrom, chromsize, dstpath, scale):
    # merge4-allsample.bw chr1 89026991 intervals ~50%
    # better to just use dense array than sparse array
    a = N.zeros(chromsize)
    for fpath in bwfiles:
        with open(fpath,mode='rb') as fobj:
            bw = BigWigFile(fobj)
            it = bw.get(chrom, 0, chromsize)
            if it is not None:
                for s,e,v in it:
                    a[s:e] += v
    #a = a/float(len(bwfiles))
    if scale is not None:
        a = a*scale
    a = N.array(a, dtype=N.float32)
    cybw.array2wiggle_chr(a, chrom, dstpath)
    return (chrom, dstpath)

def merge_bigwigs_mp(bwfiles, genome, dstpath, scale=None, np=7):
    chroms = UT.chroms(genome)
    chromfile = UT.chromsizes(genome)
    chromsizes = UT.df2dict(UT.chromdf(genome), 'chr', 'size')
    args = [(bwfiles, c, chromsizes[c], dstpath+'.{0}.wig'.format(c), scale) for c in chroms]

    rslts = UT.process_mp(merge_bigwigs_chr, args, np, doreduce=False)

    dic = dict(rslts)
    LOG.debug('concatenating chromosomes...')
    wigpath = dstpath+'.wig'
    UT.makedirs(os.path.dirname(wigpath))
    with open(wigpath, 'wb') as dst:
        for c in chroms:
            with open(dic[c],'rb') as src:
                shutil.copyfileobj(src, dst)

    LOG.debug('converting wiggle to bigwig')
    BT.wig2bw(wigpath, chromfile, dstpath)

    # clean up
    for c in chroms:
        f = dstpath+'.{0}.wig'.format(c)
        if os.path.exists(f):
            os.unlink(f)
    if os.path.exists(wigpath):
        os.unlink(wigpath)

# def array2wiggle_chr(a, chrom, dstpath):
    # possibly Cythonify
    # def _gen():
    #     i = 0
    #     # skip initial 0
    #     while(a[i]==0):
    #         i+=1
    #     st = i
    #     c = a[st] # initial non-zero
    #     i+=1
    #     while(i<len(a)):
    #         # advance until change
    #         while(i<len(a) and a[i]==c):
    #             i+=1
    #         if c!=0:
    #             # 0-based => 1-based
    #             yield '{0}\t{1}\t{2}\t{3}\n'.format(chrom,st,i,c)
    #         if i<len(a):
    #             st = i
    #             c = a[st]
    # with open(dstpath,'w') as fobj:
    #     if N.sum(a)>0: # some elements are not zero
    #         cnt = 0
    #         txt = []
    #         for line in _gen():
    #             txt.append(line)
    #             cnt += 1
    #             if cnt == 100000:
    #                 fobj.write(''.join(txt))
    #                 cnt = 0
    #                 txt = []
    #         fobj.write(txt)
    # return dstpath

<commit_msg>fix typo<commit_after>"""

.. module:: bigwig
    :synopsis: BIGWIG file related stuffs

..  moduleauthor:: Ken Sugino <ken.sugino@gmail.com>

"""

# system
import os
import subprocess
import time
import multiprocessing
from operator import iadd
import logging
logging.basicConfig(level=logging.DEBUG)
LOG = logging.getLogger(__name__)
import shutil

# 3rd party
import pandas as PD
import numpy as N

# for reading BIGWIG file
#ngslib is faster than bx but fails when parallerized
#import bx
#from bx.bbi.bigwig_file import BigWigFile
#from ngslib import wWigIO
#bx-python is not compatible with Python3 => use modified version

from jgem import utils as UT
from jgem import bedtools as BT
import jgem.cy.bw  as cybw #import array2wiggle_chr # Cython version
from jgem.cy.bw import array2wiggle_chr
from jgem.bxbbi.bigwig_file import BigWigFile

MAXSIZE = int(300e6)  # 300Mbp bigger than chr1,chrX

# BAM to genomecov BIGWIG ##################################################
def cnt_bam(fpath):
    """
    Uses samtools to count reads in BAM file.

    Args:
        fpath (str): path to BAM file

    Returns:
        int. The number of aligned reads.

    Raises:
        RuntimeError
    """
    cache = fpath+'.flagstat'
    bai = fpath+'.bai'
    if not os.path.exists(fpath):
        raise RuntimeError('file {0} does not exist'.format(fpath))
    if not os.path.exists(bai):
        cmd = ['samtools', 'index', fpath]
        subprocess.call(cmd)
    if os.path.exists(cache) and os.path.getmtime(cache)>os.path.getmtime(fpath):
        out = open(cache,'r').read()
    else:
        cmd = ['samtools', 'flagstat', fpath]
        out = subprocess.check_output(cmd)
        open(cache,'wb').write(out)
    firstline = out.split('\n')[0].split()
    return int(firstline[0])+int(firstline[2])

def wig2bw(wigpath, chromsizes, bwpath):
    pass

def bam2bw(fpath, chromsizes, bpath, aligned=None):
    """
    Generate normalized coverage from BAM

    Args:
        fpath (str): path to BAM
        chromsizes (str): path to chromsizes file
        bpath (str): path to BIGWIG
        aligned (int): number of aligned reads, if None uses samtools to find it from BAM

    Requires Bedtools (genomeCoverageBed) and Kent Tool (wigToBigWig)

    """
    # countreads
    if aligned is None:
        aligned =   cnt_bam(fpath)
    scale = 1000000./float(aligned)
    # convert_to_wig
    tpath = bpath +'.wig'
    UT.makedirs(os.path.dirname(tpath))
    tfobj = open(tpath,'wb')
    cmd1 = ['genomeCoverageBed', '-split', '-bg', '-ibam', fpath, '-g', chromsizes, '-scale', str(scale)]
    p1 = subprocess.Popen(cmd1, stdout=tfobj)
    p1.wait()
    tfobj.close()

    # convet_wig_to_bigwig
    cmd2 = ['wigToBigWig', tpath, chromsizes, bpath]
    p2 = subprocess.call(cmd2)

    # remove_temporary_file
    os.remove(tpath)

# bw2bed based on #########################################################
# https://github.com/CGATOxford/cgat/blob/master/scripts/wig2bed.py
# [TODO] report bug to AndreasHeger line 80 "elif => if"

def block_iter(infile, chrom, chunk=int(10e6)):
    "BigWig file iterator"
    with open(infile, mode='rb') as fobj:
        bw = BigWigFile(fobj)
        for x in range(0,MAXSIZE,chunk): # 10Mbp chunk
            iterator = bw.get(chrom, x, x+chunk)
            if iterator is None:
                raise StopIteration
            for v in iterator:
                yield v

# def block_iter_ngs(infile, chrom, chunk=int(10e6)):
#     # ngslib is faster than bx but fails when parallerized
#     wWigIO.open(infile)
#     for x in range(0,MAXSIZE,chunk): # 10Mbp chunk
#         iterator = wWigIO.getIntervals(infile, chrom, x, x+chunk)
#         if not iterator:
#             raise StopIteration
#         for v in iterator:
#             yield v
#     wWigIO.close(infile)

def apply_threshold(infile, threshold, chroms):
    """Only returns intervals exceeding the threshold

    Args:
        infile: path to bigwig
        threshold: positive float
        chroms: list of chromosome names

    Yields:
        intervals (chr,st,ed)
    """
    for chrom in chroms:
        last_start, last_end = -1, 0
        #for start, end, value in block_iter_ngs(infile, chrom):
        for start, end, value in block_iter(infile, chrom):
            # start=>end has value
            d = start - last_end
            # d: distance from last_end
            # d> 0 => last_end==>start value is 0 assume threshold>=0
            if (d > 0 or value <= threshold):
                if last_start >= 0: # if there's un-yielded interval then yield
                    yield chrom, last_start, last_end
                last_start = -1 # reset (no hit)
            #elif last_start < 0 and value > threshold:
            # this part is a bug from the original code
            # original code will skip the next interval
            if last_start < 0 and value > threshold:
                # last_start <0 => no current interval
                # and value above threshod ==> start new interval
                last_start = start
            last_end = end
        if last_start >= 0:
            yield chrom, last_start, end

def bw2bed(bwfile, bedfile, chroms, th, compress=True):
    """Transform BigWig genomeCov to binary BED by thresholding.
    Makes result file (bwfile[:-3]+'.binary%g.bed'.format(th))

    Args:
        bwfile: path to BigWig file
        chroms: list of chromosome names
        th: coverage threshold

    Returns:
        path to generated BED file
    """
    bedbase = bedfile[:-3] if bedfile[-3:]=='.gz' else bedfile
    #bedfile = '{0}.binary{1:g}.bed'.format(bwfile[:-3], th)
    if UT.notstale(bwfile, bedbase+'.gz'):
        return bedbase+'.gz'
    # make sure bwfile exists
    if not ( os.path.exists(bwfile) ):
        raise RuntimeError('BigWig file {0} does not exist.'.format(bwfile))
    processor = apply_threshold(bwfile,th, chroms)
    UT.makedirs(os.path.dirname(bedfile))
    out = open(bedbase,'w')
    out.write(''.join(['%s\t%i\t%i\n' % x for x in processor]))
    #out.write('\n') #<= this introduces space inbetween chroms in mp ode
    # which terminates bedtools at chr1
    out.close()
    if compress:
        return UT.compress(bedbase)
    return bedbase

def bw2bed_mp(bwfile, bedfile, chroms, th, np=4):
    """ multi CPU version of bw2bed """

    args = []
    files = []
    for chrom in chroms:
        bedchromfile = bedfile+'.{0}.bed.gz'.format(chrom)
        files.append(bedchromfile)
        args.append((bwfile,bedchromfile,[chrom],th,False))

    rslts = UT.process_mp(bw2bed, args, np=np, doreduce=False)

    # concatenate gz files
    bedbase = bedfile[:-3] if bedfile[-3:]=='.gz' else bedfile
    with open(bedbase, 'wb') as dst:
        for f in rslts:
            with open(f, 'rb') as src:
                shutil.copyfileobj(src, dst)
    # !!! bedtool gzip problem againg !!!
    # bedtools only process first one if just concatenate gzipped files
    # => concatenate unzipped and gzip whole thing at the end
    bedfile = UT.compress(bedbase)

    # clean up temp files
    for f in rslts:
        os.unlink(f)

    return bedfile

### covbp, totbp, avgcov from bigwig array #################################

def get_totbp_covbp_bw(bwfile, genome, chroms=None):
    """ Calculate total bp, covered bp, mean coverage, covered %.

    Args:
        bwfile: bigwig file
        genome: UCSC genome name
        chroms (list): of chromosomes

    Returns:
        Pandas dataframe

    """
    chromdf = UT.chromdf(genome).set_index('chr')['size']
    def one(chrom):
        csize = chromdf.ix[chrom]
        a = get_bigwig_as_array(bwfile, chrom, 0, csize)
        totbp = N.sum(a)
        covbp = N.sum(a>0)
        acov = float(totbp)/covbp
        covp = float(covbp)/csize
        return {'totbp':totbp,'covbp':covbp,'acov':acov,'cov%':covp}
    if chroms is None:
        chroms = UT.chroms(genome)
    df = PD.DataFrame({x: one(x) for x in chroms})
    return df

### Merge BigWigs ##########################################################

def get_bigwig_as_array(bwfile, chrom, st, ed):
    """Get BIGWIG coverage values as array of size (ed-st). Array start corresponds to st.
    0-based

    Args:
        bwfile: path to BIGWIG
        chrom (str): chromosome name
        st (int): start position
        ed (int): end position

    Returns:
        Numpy array of size (ed-st)
    """
    # with open(bwfile, mode='rb') as fobj:
    #     bw = BigWigFile(fobj)
    #     it = bw.get(chrom,st,ed)
    #     a = N.zeros(ed-st)
    #     for s,e,v in it:
    #         a[s-st:e-st] += v
    # return a
    with open(bwfile, mode='rb') as fobj:
        bw = BigWigFile(fobj)
        a = bw.get_as_array(chrom,st,ed)
        a[N.isnan(a)]=0.
    return a

def merge_bigwigs_chr(bwfiles, chrom, chromsize, dstpath, scale):
    # merge4-allsample.bw chr1 89026991 intervals ~50%
    # better to just use dense array than sparse array
    a = N.zeros(chromsize)
    for fpath in bwfiles:
        with open(fpath,mode='rb') as fobj:
            bw = BigWigFile(fobj)
            it = bw.get(chrom, 0, chromsize)
            if it is not None:
                for s,e,v in it:
                    a[s:e] += v
    #a = a/float(len(bwfiles))
    if scale is not None:
        a = a*scale
    a = N.array(a, dtype=N.float32)
    cybw.array2wiggle_chr(a, chrom, dstpath)
    return (chrom, dstpath)

def merge_bigwigs_mp(bwfiles, genome, dstpath, scale=None, np=7):
    chroms = UT.chroms(genome)
    chromfile = UT.chromsizes(genome)
    chromsizes = UT.df2dict(UT.chromdf(genome), 'chr', 'size')
    args = [(bwfiles, c, chromsizes[c], dstpath+'.{0}.wig'.format(c), scale) for c in chroms]

    rslts = UT.process_mp(merge_bigwigs_chr, args, np, doreduce=False)

    dic = dict(rslts)
    LOG.debug('concatenating chromosomes...')
    wigpath = dstpath+'.wig'
    UT.makedirs(os.path.dirname(wigpath))
    with open(wigpath, 'wb') as dst:
        for c in chroms:
            with open(dic[c],'rb') as src:
                shutil.copyfileobj(src, dst)

    LOG.debug('converting wiggle to bigwig')
    BT.wig2bw(wigpath, chromfile, dstpath)

    # clean up
    for c in chroms:
        f = dstpath+'.{0}.wig'.format(c)
        if os.path.exists(f):
            os.unlink(f)
    if os.path.exists(wigpath):
        os.unlink(wigpath)

# def array2wiggle_chr(a, chrom, dstpath):
    # possibly Cythonify
    # def _gen():
    #     i = 0
    #     # skip initial 0
    #     while(a[i]==0):
    #         i+=1
    #     st = i
    #     c = a[st] # initial non-zero
    #     i+=1
    #     while(i<len(a)):
    #         # advance until change
    #         while(i<len(a) and a[i]==c):
    #             i+=1
    #         if c!=0:
    #             # 0-based => 1-based
    #             yield '{0}\t{1}\t{2}\t{3}\n'.format(chrom,st,i,c)
    #         if i<len(a):
    #             st = i
    #             c = a[st]
    # with open(dstpath,'w') as fobj:
    #     if N.sum(a)>0: # some elements are not zero
    #         cnt = 0
    #         txt = []
    #         for line in _gen():
    #             txt.append(line)
    #             cnt += 1
    #             if cnt == 100000:
    #                 fobj.write(''.join(txt))
    #                 cnt = 0
    #                 txt = []
    #         fobj.write(txt)
    # return dstpath

<|endoftext|>