<commit_before>import re
import falcon
import ipaddress
import mimetypes
import os
import json
import types
import click
from time import sleep
from certidude.wrappers import Request, Certificate, CertificateAuthorityConfig
from certidude.auth import login_required
from pyasn1.codec.der import decoder
from datetime import datetime, date
from jinja2 import Environment, PackageLoader, Template

env = Environment(loader=PackageLoader("certidude", "templates"))

RE_HOSTNAME = "^(([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])\.)*([A-Za-z0-9]|[A-Za-z0-9][A-Za-z0-9\-]*[A-Za-z0-9])$"

def omit(**kwargs):
    return dict([(key,value) for (key, value) in kwargs.items() if value])

def authorize_admin(func):
    def wrapped(self, req, resp, *args, **kwargs):
        authority = kwargs.get("ca")

        # Parse remote IPv4/IPv6 address
        remote_addr = ipaddress.ip_network(req.env["REMOTE_ADDR"])

        # Check for administration subnet whitelist
        print("Comparing:", authority.admin_subnets, "To:", remote_addr)
        for subnet in authority.admin_subnets:
            if subnet.overlaps(remote_addr):
                break
        else:
            raise falcon.HTTPForbidden("Forbidden", "Remote address %s not whitelisted" % remote_addr)

        # Check for username whitelist
        kerberos_username, kerberos_realm = kwargs.get("user")
        if kerberos_username not in authority.admin_users:
            raise falcon.HTTPForbidden("Forbidden", "User %s not whitelisted" % kerberos_username)

        # Retain username, TODO: Better abstraction with username, e-mail, sn, gn?
        kwargs["user"] = kerberos_username
        return func(self, req, resp, *args, **kwargs)
    return wrapped

def pop_certificate_authority(func):
    def wrapped(self, req, resp, *args, **kwargs):
        kwargs["ca"] = self.config.instantiate_authority(kwargs["ca"])
        return func(self, req, resp, *args, **kwargs)
    return wrapped

def validate_common_name(func):
    def wrapped(*args, **kwargs):
        if not re.match(RE_HOSTNAME, kwargs["cn"]):
            raise falcon.HTTPBadRequest("Invalid CN", "Common name supplied with request didn't pass the validation regex")
        return func(*args, **kwargs)
    return wrapped

class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
        if isinstance(obj, date):
            return obj.strftime("%Y-%m-%d")
        if isinstance(obj, map):
            return tuple(obj)
        if isinstance(obj, types.GeneratorType):
            return tuple(obj)
        return json.JSONEncoder.default(self, obj)

def serialize(func):
    """
    Falcon response serialization
    """
    def wrapped(instance, req, resp, **kwargs):
        assert not req.get_param("unicode") or req.get_param("unicode") == u"✓", "Unicode sanity check failed"
        resp.set_header("Cache-Control", "no-cache, no-store, must-revalidate");
        resp.set_header("Pragma", "no-cache");
        resp.set_header("Expires", "0");
        r = func(instance, req, resp, **kwargs)
        if not resp.body:
            if not req.client_accepts_json:
                raise falcon.HTTPUnsupportedMediaType(
                    "This API only supports the JSON media type.",
                    href="http://docs.examples.com/api/json")
            resp.set_header("Content-Type", "application/json")
            resp.body = json.dumps(r, cls=MyEncoder)
        return r
    return wrapped

def templatize(path):
    template = env.get_template(path)
    def wrapper(func):
        def wrapped(instance, req, resp, *args, **kwargs):
            assert not req.get_param("unicode") or req.get_param("unicode") == u"✓", "Unicode sanity check failed"
            r = func(instance, req, resp, *args, **kwargs)
            r.pop("self")
            if not resp.body:
                if  req.get_header("Accept") == "application/json":
                    resp.set_header("Cache-Control", "no-cache, no-store, must-revalidate");
                    resp.set_header("Pragma", "no-cache");
                    resp.set_header("Expires", "0");
                    resp.set_header("Content-Type", "application/json")
                    resp.body = json.dumps(r, cls=MyEncoder)
                    return r
                else:
                    resp.set_header("Content-Type", "text/html")
                    resp.body = template.render(request=req, **r)
                    return r
        return wrapped
    return wrapper

class CertificateAuthorityBase(object):
    def __init__(self, config):
        self.config = config

class RevocationListResource(CertificateAuthorityBase):
    @pop_certificate_authority
    def on_get(self, req, resp, ca):
        resp.set_header("Content-Type", "application/x-pkcs7-crl")
        resp.append_header("Content-Disposition", "attachment; filename=%s.crl" % ca.slug)
        resp.body = ca.export_crl()

class SignedCertificateDetailResource(CertificateAuthorityBase):
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca, cn):
        path = os.path.join(ca.signed_dir, cn + ".pem")
        if not os.path.exists(path):
            raise falcon.HTTPNotFound()
        resp.stream = open(path, "rb")
        resp.append_header("Content-Disposition", "attachment; filename=%s.crt" % cn)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    @validate_common_name
    def on_delete(self, req, resp, ca, cn, user):
        ca.revoke(cn)

class SignedCertificateListResource(CertificateAuthorityBase):
    @serialize
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca):
        for j in authority.get_signed():
            yield omit(
                key_type=j.key_type,
                key_length=j.key_length,
                subject=j.distinguished_name,
                cn=j.common_name,
                c=j.country_code,
                st=j.state_or_county,
                l=j.city,
                o=j.organization,
                ou=j.organizational_unit,
                fingerprint=j.fingerprint)

class RequestDetailResource(CertificateAuthorityBase):
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca, cn):
        """
        Fetch certificate signing request as PEM
        """
        path = os.path.join(ca.request_dir, cn + ".pem")
        if not os.path.exists(path):
            raise falcon.HTTPNotFound()
        resp.stream = open(path, "rb")
        resp.append_header("Content-Type", "application/x-x509-user-cert")
        resp.append_header("Content-Disposition", "attachment; filename=%s.csr" % cn)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    @validate_common_name
    def on_patch(self, req, resp, ca, cn, user):
        """
        Sign a certificate signing request
        """
        csr = ca.get_request(cn)
        cert = ca.sign(csr, overwrite=True, delete=True)
        os.unlink(csr.path)
        resp.body = "Certificate successfully signed"
        resp.status = falcon.HTTP_201
        resp.location = os.path.join(req.relative_uri, "..", "..", "signed", cn)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    def on_delete(self, req, resp, ca, cn, user):
        ca.delete_request(cn)

class RequestListResource(CertificateAuthorityBase):
    @serialize
    @pop_certificate_authority
    def on_get(self, req, resp, ca):
        for j in ca.get_requests():
            yield omit(
                key_type=j.key_type,
                key_length=j.key_length,
                subject=j.distinguished_name,
                cn=j.common_name,
                c=j.country_code,
                st=j.state_or_county,
                l=j.city,
                o=j.organization,
                ou=j.organizational_unit,
                fingerprint=j.fingerprint())

    @pop_certificate_authority
    def on_post(self, req, resp, ca):
        """
        Submit certificate signing request (CSR) in PEM format
        """
        # Parse remote IPv4/IPv6 address
        remote_addr = ipaddress.ip_network(req.env["REMOTE_ADDR"])

        # Check for CSR submission whitelist
        if ca.request_subnets:
            for subnet in ca.request_subnets:
                if subnet.overlaps(remote_addr):
                    break
            else:
               raise falcon.HTTPForbidden("Forbidden", "IP address %s not whitelisted" % remote_addr)

        if req.get_header("Content-Type") != "application/pkcs10":
            raise falcon.HTTPUnsupportedMediaType(
                "This API call accepts only application/pkcs10 content type")

        body = req.stream.read(req.content_length)
        csr = Request(body)

        # Check if this request has been already signed and return corresponding certificte if it has been signed
        try:
            cert_buf = ca.get_certificate(csr.common_name)
        except FileNotFoundError:
            pass
        else:
            cert = Certificate(cert_buf)
            if cert.pubkey == csr.pubkey:
                resp.status = falcon.HTTP_SEE_OTHER
                resp.location = os.path.join(os.path.dirname(req.relative_uri), "signed", csr.common_name)
                return

        # TODO: check for revoked certificates and return HTTP 410 Gone

        # Process automatic signing if the IP address is whitelisted and autosigning was requested
        if req.get_param("autosign").lower() in ("yes", "1", "true"):
            for subnet in ca.autosign_subnets:
                if subnet.overlaps(remote_addr):
                    try:
                        resp.append_header("Content-Type", "application/x-x509-user-cert")
                        resp.body = ca.sign(csr).dump()
                        return
                    except FileExistsError: # Certificate already exists, try to save the request
                        pass
                    break

        # Attempt to save the request otherwise
        try:
            request = ca.store_request(body)
        except FileExistsError:
            raise falcon.HTTPConflict(
                "CSR with such CN already exists",
                "Will not overwrite existing certificate signing request, explicitly delete CSR and try again")

        # Wait the certificate to be signed if waiting is requested
        if req.get_param("wait"):
            url_template = os.getenv("PUSH_SUBSCRIBE")
            if url_template:
                # Redirect to nginx pub/sub
                url = url_template % dict(channel=request.fingerprint())
                click.echo("Redirecting to: %s"  % url)
                resp.status = falcon.HTTP_SEE_OTHER
                resp.append_header("Location", url)
            else:
                click.echo("Using dummy streaming mode, please switch to nginx in production!", err=True)
                # Dummy streaming mode
                while True:
                    sleep(1)
                    if not ca.request_exists(csr.common_name):
                        resp.append_header("Content-Type", "application/x-x509-user-cert")
                        resp.status = falcon.HTTP_201 # Certificate was created
                        resp.body = ca.get_certificate(csr.common_name)
                        break
        else:
            # Request was accepted, but not processed
            resp.status = falcon.HTTP_202

class CertificateStatusResource(CertificateAuthorityBase):
    """
    openssl ocsp -issuer CAcert_class1.pem -serial 0x<serial no in hex> -url http://localhost -CAfile cacert_both.pem
    """
    def on_post(self, req, resp, ca):
        ocsp_request = req.stream.read(req.content_length)
        for component in decoder.decode(ocsp_request):
            click.echo(component)
        resp.append_header("Content-Type", "application/ocsp-response")
        resp.status = falcon.HTTP_200
        raise NotImplementedError()

class CertificateAuthorityResource(CertificateAuthorityBase):
    @pop_certificate_authority
    def on_get(self, req, resp, ca):
        path = os.path.join(ca.certificate.path)
        resp.stream = open(path, "rb")
        resp.append_header("Content-Disposition", "attachment; filename=%s.crt" % ca.slug)

class IndexResource(CertificateAuthorityBase):
    @login_required
    @pop_certificate_authority
    @authorize_admin
    @templatize("index.html")
    def on_get(self, req, resp, ca, user):
        return locals()

class ApplicationConfigurationResource(CertificateAuthorityBase):
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca, cn):
        ctx = dict(
            cn = cn,
            certificate = ca.get_certificate(cn),
            ca_certificate = open(ca.certificate.path, "r").read())
        resp.append_header("Content-Type", "application/ovpn")
        resp.append_header("Content-Disposition", "attachment; filename=%s.ovpn" % cn)
        resp.body = Template(open("/etc/openvpn/%s.template" % ca.slug).read()).render(ctx)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    @validate_common_name
    def on_put(self, req, resp, user, ca, cn=None):
        pkey_buf, req_buf, cert_buf = ca.create_bundle(cn)

        ctx = dict(
            private_key = pkey_buf,
            certificate = cert_buf,
            ca_certificate = ca.certificate.dump())

        resp.append_header("Content-Type", "application/ovpn")
        resp.append_header("Content-Disposition", "attachment; filename=%s.ovpn" % cn)
        resp.body = Template(open("/etc/openvpn/%s.template" % ca.slug).read()).render(ctx)

class StaticResource(object):
    def __init__(self, root):
        self.root = os.path.realpath(root)

    def __call__(self, req, resp):

        path = os.path.realpath(os.path.join(self.root, req.path[1:]))
        if not path.startswith(self.root):
            raise falcon.HTTPForbidden

        print("Serving:", path)
        if os.path.exists(path):
            content_type, content_encoding = mimetypes.guess_type(path)
            if content_type:
                resp.append_header("Content-Type", content_type)
            if content_encoding:
                resp.append_header("Content-Encoding", content_encoding)
            resp.append_header("Content-Disposition", "attachment")
            resp.stream = open(path, "rb")
        else:
            resp.status = falcon.HTTP_404
            resp.body = "File '%s' not found" % req.path

def certidude_app():
    config = CertificateAuthorityConfig()

    app = falcon.API()
    app.add_route("/api/{ca}/ocsp/", CertificateStatusResource(config))
    app.add_route("/api/{ca}/signed/{cn}/openvpn", ApplicationConfigurationResource(config))
    app.add_route("/api/{ca}/certificate/", CertificateAuthorityResource(config))
    app.add_route("/api/{ca}/revoked/", RevocationListResource(config))
    app.add_route("/api/{ca}/signed/{cn}/", SignedCertificateDetailResource(config))
    app.add_route("/api/{ca}/signed/", SignedCertificateListResource(config))
    app.add_route("/api/{ca}/request/{cn}/", RequestDetailResource(config))
    app.add_route("/api/{ca}/request/", RequestListResource(config))
    app.add_route("/api/{ca}/", IndexResource(config))

    return app
<commit_msg>api: Fixed CSR processing if autosign GET variable was not present<commit_after>import re
import falcon
import ipaddress
import mimetypes
import os
import json
import types
import click
from time import sleep
from certidude.wrappers import Request, Certificate, CertificateAuthorityConfig
from certidude.auth import login_required
from pyasn1.codec.der import decoder
from datetime import datetime, date
from jinja2 import Environment, PackageLoader, Template

env = Environment(loader=PackageLoader("certidude", "templates"))

RE_HOSTNAME = "^(([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])\.)*([A-Za-z0-9]|[A-Za-z0-9][A-Za-z0-9\-]*[A-Za-z0-9])$"

def omit(**kwargs):
    return dict([(key,value) for (key, value) in kwargs.items() if value])

def authorize_admin(func):
    def wrapped(self, req, resp, *args, **kwargs):
        authority = kwargs.get("ca")

        # Parse remote IPv4/IPv6 address
        remote_addr = ipaddress.ip_network(req.env["REMOTE_ADDR"])

        # Check for administration subnet whitelist
        print("Comparing:", authority.admin_subnets, "To:", remote_addr)
        for subnet in authority.admin_subnets:
            if subnet.overlaps(remote_addr):
                break
        else:
            raise falcon.HTTPForbidden("Forbidden", "Remote address %s not whitelisted" % remote_addr)

        # Check for username whitelist
        kerberos_username, kerberos_realm = kwargs.get("user")
        if kerberos_username not in authority.admin_users:
            raise falcon.HTTPForbidden("Forbidden", "User %s not whitelisted" % kerberos_username)

        # Retain username, TODO: Better abstraction with username, e-mail, sn, gn?
        kwargs["user"] = kerberos_username
        return func(self, req, resp, *args, **kwargs)
    return wrapped

def pop_certificate_authority(func):
    def wrapped(self, req, resp, *args, **kwargs):
        kwargs["ca"] = self.config.instantiate_authority(kwargs["ca"])
        return func(self, req, resp, *args, **kwargs)
    return wrapped

def validate_common_name(func):
    def wrapped(*args, **kwargs):
        if not re.match(RE_HOSTNAME, kwargs["cn"]):
            raise falcon.HTTPBadRequest("Invalid CN", "Common name supplied with request didn't pass the validation regex")
        return func(*args, **kwargs)
    return wrapped

class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
        if isinstance(obj, date):
            return obj.strftime("%Y-%m-%d")
        if isinstance(obj, map):
            return tuple(obj)
        if isinstance(obj, types.GeneratorType):
            return tuple(obj)
        return json.JSONEncoder.default(self, obj)

def serialize(func):
    """
    Falcon response serialization
    """
    def wrapped(instance, req, resp, **kwargs):
        assert not req.get_param("unicode") or req.get_param("unicode") == u"✓", "Unicode sanity check failed"
        resp.set_header("Cache-Control", "no-cache, no-store, must-revalidate");
        resp.set_header("Pragma", "no-cache");
        resp.set_header("Expires", "0");
        r = func(instance, req, resp, **kwargs)
        if not resp.body:
            if not req.client_accepts_json:
                raise falcon.HTTPUnsupportedMediaType(
                    "This API only supports the JSON media type.",
                    href="http://docs.examples.com/api/json")
            resp.set_header("Content-Type", "application/json")
            resp.body = json.dumps(r, cls=MyEncoder)
        return r
    return wrapped

def templatize(path):
    template = env.get_template(path)
    def wrapper(func):
        def wrapped(instance, req, resp, *args, **kwargs):
            assert not req.get_param("unicode") or req.get_param("unicode") == u"✓", "Unicode sanity check failed"
            r = func(instance, req, resp, *args, **kwargs)
            r.pop("self")
            if not resp.body:
                if  req.get_header("Accept") == "application/json":
                    resp.set_header("Cache-Control", "no-cache, no-store, must-revalidate");
                    resp.set_header("Pragma", "no-cache");
                    resp.set_header("Expires", "0");
                    resp.set_header("Content-Type", "application/json")
                    resp.body = json.dumps(r, cls=MyEncoder)
                    return r
                else:
                    resp.set_header("Content-Type", "text/html")
                    resp.body = template.render(request=req, **r)
                    return r
        return wrapped
    return wrapper

class CertificateAuthorityBase(object):
    def __init__(self, config):
        self.config = config

class RevocationListResource(CertificateAuthorityBase):
    @pop_certificate_authority
    def on_get(self, req, resp, ca):
        resp.set_header("Content-Type", "application/x-pkcs7-crl")
        resp.append_header("Content-Disposition", "attachment; filename=%s.crl" % ca.slug)
        resp.body = ca.export_crl()

class SignedCertificateDetailResource(CertificateAuthorityBase):
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca, cn):
        path = os.path.join(ca.signed_dir, cn + ".pem")
        if not os.path.exists(path):
            raise falcon.HTTPNotFound()
        resp.stream = open(path, "rb")
        resp.append_header("Content-Disposition", "attachment; filename=%s.crt" % cn)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    @validate_common_name
    def on_delete(self, req, resp, ca, cn, user):
        ca.revoke(cn)

class SignedCertificateListResource(CertificateAuthorityBase):
    @serialize
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca):
        for j in authority.get_signed():
            yield omit(
                key_type=j.key_type,
                key_length=j.key_length,
                subject=j.distinguished_name,
                cn=j.common_name,
                c=j.country_code,
                st=j.state_or_county,
                l=j.city,
                o=j.organization,
                ou=j.organizational_unit,
                fingerprint=j.fingerprint)

class RequestDetailResource(CertificateAuthorityBase):
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca, cn):
        """
        Fetch certificate signing request as PEM
        """
        path = os.path.join(ca.request_dir, cn + ".pem")
        if not os.path.exists(path):
            raise falcon.HTTPNotFound()
        resp.stream = open(path, "rb")
        resp.append_header("Content-Type", "application/x-x509-user-cert")
        resp.append_header("Content-Disposition", "attachment; filename=%s.csr" % cn)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    @validate_common_name
    def on_patch(self, req, resp, ca, cn, user):
        """
        Sign a certificate signing request
        """
        csr = ca.get_request(cn)
        cert = ca.sign(csr, overwrite=True, delete=True)
        os.unlink(csr.path)
        resp.body = "Certificate successfully signed"
        resp.status = falcon.HTTP_201
        resp.location = os.path.join(req.relative_uri, "..", "..", "signed", cn)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    def on_delete(self, req, resp, ca, cn, user):
        ca.delete_request(cn)

class RequestListResource(CertificateAuthorityBase):
    @serialize
    @pop_certificate_authority
    def on_get(self, req, resp, ca):
        for j in ca.get_requests():
            yield omit(
                key_type=j.key_type,
                key_length=j.key_length,
                subject=j.distinguished_name,
                cn=j.common_name,
                c=j.country_code,
                st=j.state_or_county,
                l=j.city,
                o=j.organization,
                ou=j.organizational_unit,
                fingerprint=j.fingerprint())

    @pop_certificate_authority
    def on_post(self, req, resp, ca):
        """
        Submit certificate signing request (CSR) in PEM format
        """
        # Parse remote IPv4/IPv6 address
        remote_addr = ipaddress.ip_network(req.env["REMOTE_ADDR"])

        # Check for CSR submission whitelist
        if ca.request_subnets:
            for subnet in ca.request_subnets:
                if subnet.overlaps(remote_addr):
                    break
            else:
               raise falcon.HTTPForbidden("Forbidden", "IP address %s not whitelisted" % remote_addr)

        if req.get_header("Content-Type") != "application/pkcs10":
            raise falcon.HTTPUnsupportedMediaType(
                "This API call accepts only application/pkcs10 content type")

        body = req.stream.read(req.content_length)
        csr = Request(body)

        # Check if this request has been already signed and return corresponding certificte if it has been signed
        try:
            cert_buf = ca.get_certificate(csr.common_name)
        except FileNotFoundError:
            pass
        else:
            cert = Certificate(cert_buf)
            if cert.pubkey == csr.pubkey:
                resp.status = falcon.HTTP_SEE_OTHER
                resp.location = os.path.join(os.path.dirname(req.relative_uri), "signed", csr.common_name)
                return

        # TODO: check for revoked certificates and return HTTP 410 Gone

        # Process automatic signing if the IP address is whitelisted and autosigning was requested
        if req.get_param("autosign") in ("yes", "1", "true"):
            for subnet in ca.autosign_subnets:
                if subnet.overlaps(remote_addr):
                    try:
                        resp.append_header("Content-Type", "application/x-x509-user-cert")
                        resp.body = ca.sign(csr).dump()
                        return
                    except FileExistsError: # Certificate already exists, try to save the request
                        pass
                    break

        # Attempt to save the request otherwise
        try:
            request = ca.store_request(body)
        except FileExistsError:
            raise falcon.HTTPConflict(
                "CSR with such CN already exists",
                "Will not overwrite existing certificate signing request, explicitly delete CSR and try again")

        # Wait the certificate to be signed if waiting is requested
        if req.get_param("wait"):
            url_template = os.getenv("PUSH_SUBSCRIBE")
            if url_template:
                # Redirect to nginx pub/sub
                url = url_template % dict(channel=request.fingerprint())
                click.echo("Redirecting to: %s"  % url)
                resp.status = falcon.HTTP_SEE_OTHER
                resp.append_header("Location", url)
            else:
                click.echo("Using dummy streaming mode, please switch to nginx in production!", err=True)
                # Dummy streaming mode
                while True:
                    sleep(1)
                    if not ca.request_exists(csr.common_name):
                        resp.append_header("Content-Type", "application/x-x509-user-cert")
                        resp.status = falcon.HTTP_201 # Certificate was created
                        resp.body = ca.get_certificate(csr.common_name)
                        break
        else:
            # Request was accepted, but not processed
            resp.status = falcon.HTTP_202

class CertificateStatusResource(CertificateAuthorityBase):
    """
    openssl ocsp -issuer CAcert_class1.pem -serial 0x<serial no in hex> -url http://localhost -CAfile cacert_both.pem
    """
    def on_post(self, req, resp, ca):
        ocsp_request = req.stream.read(req.content_length)
        for component in decoder.decode(ocsp_request):
            click.echo(component)
        resp.append_header("Content-Type", "application/ocsp-response")
        resp.status = falcon.HTTP_200
        raise NotImplementedError()

class CertificateAuthorityResource(CertificateAuthorityBase):
    @pop_certificate_authority
    def on_get(self, req, resp, ca):
        path = os.path.join(ca.certificate.path)
        resp.stream = open(path, "rb")
        resp.append_header("Content-Disposition", "attachment; filename=%s.crt" % ca.slug)

class IndexResource(CertificateAuthorityBase):
    @login_required
    @pop_certificate_authority
    @authorize_admin
    @templatize("index.html")
    def on_get(self, req, resp, ca, user):
        return locals()

class ApplicationConfigurationResource(CertificateAuthorityBase):
    @pop_certificate_authority
    @validate_common_name
    def on_get(self, req, resp, ca, cn):
        ctx = dict(
            cn = cn,
            certificate = ca.get_certificate(cn),
            ca_certificate = open(ca.certificate.path, "r").read())
        resp.append_header("Content-Type", "application/ovpn")
        resp.append_header("Content-Disposition", "attachment; filename=%s.ovpn" % cn)
        resp.body = Template(open("/etc/openvpn/%s.template" % ca.slug).read()).render(ctx)

    @login_required
    @pop_certificate_authority
    @authorize_admin
    @validate_common_name
    def on_put(self, req, resp, user, ca, cn=None):
        pkey_buf, req_buf, cert_buf = ca.create_bundle(cn)

        ctx = dict(
            private_key = pkey_buf,
            certificate = cert_buf,
            ca_certificate = ca.certificate.dump())

        resp.append_header("Content-Type", "application/ovpn")
        resp.append_header("Content-Disposition", "attachment; filename=%s.ovpn" % cn)
        resp.body = Template(open("/etc/openvpn/%s.template" % ca.slug).read()).render(ctx)

class StaticResource(object):
    def __init__(self, root):
        self.root = os.path.realpath(root)

    def __call__(self, req, resp):

        path = os.path.realpath(os.path.join(self.root, req.path[1:]))
        if not path.startswith(self.root):
            raise falcon.HTTPForbidden

        print("Serving:", path)
        if os.path.exists(path):
            content_type, content_encoding = mimetypes.guess_type(path)
            if content_type:
                resp.append_header("Content-Type", content_type)
            if content_encoding:
                resp.append_header("Content-Encoding", content_encoding)
            resp.append_header("Content-Disposition", "attachment")
            resp.stream = open(path, "rb")
        else:
            resp.status = falcon.HTTP_404
            resp.body = "File '%s' not found" % req.path

def certidude_app():
    config = CertificateAuthorityConfig()

    app = falcon.API()
    app.add_route("/api/{ca}/ocsp/", CertificateStatusResource(config))
    app.add_route("/api/{ca}/signed/{cn}/openvpn", ApplicationConfigurationResource(config))
    app.add_route("/api/{ca}/certificate/", CertificateAuthorityResource(config))
    app.add_route("/api/{ca}/revoked/", RevocationListResource(config))
    app.add_route("/api/{ca}/signed/{cn}/", SignedCertificateDetailResource(config))
    app.add_route("/api/{ca}/signed/", SignedCertificateListResource(config))
    app.add_route("/api/{ca}/request/{cn}/", RequestDetailResource(config))
    app.add_route("/api/{ca}/request/", RequestListResource(config))
    app.add_route("/api/{ca}/", IndexResource(config))

    return app
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#   Copyright 2017 Kaede Hoshikawa
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

__all__ = ["__version__"]

_tag_version = (0, 1, 1)

_dev = 0

_version_fragments = [str(i) for i in _tag_version[:3]]

if _dev is not None:
    _version_fragments.append(f"dev{_dev}")

__version__ = ".".join(_version_fragments)
<commit_msg>Move master version to 0.1.2.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#   Copyright 2017 Kaede Hoshikawa
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

__all__ = ["__version__"]

_tag_version = (0, 1, 2)

_dev = 0

_version_fragments = [str(i) for i in _tag_version[:3]]

if _dev is not None:
    _version_fragments.append(f"dev{_dev}")

__version__ = ".".join(_version_fragments)
<|endoftext|>


<commit_before>{
  "targets": [
    {
      'target_name': 'libmongoose',
      'type': 'static_library',
      'include_dirs': [
        'mongoose/',
      ],
      'sources': [
        'mongoose/mongoose.c',
      ],
    },
  ],
}
<commit_msg>Disable mongoose unused features.<commit_after>{
  "targets": [
    {
      'target_name': 'libmongoose',
      'type': 'static_library',
      'include_dirs': [
        'mongoose/',
      ],
      'defines': [
        'MONGOOSE_NO_AUTH',
        'MONGOOSE_NO_CGI',
        'MONGOOSE_NO_DAV',
        'MONGOOSE_NO_LOGGING',
        'MONGOOSE_NO_WEBSOCKET',
      ],
      'sources': [
        'mongoose/mongoose.c',
      ],
    },
  ],
}
<|endoftext|>


<commit_before># This file is part of Viper - https://github.com/viper-framework/viper
# See the file 'LICENSE' for copying permission.

import os
import sys
import glob
import atexit
import readline
import traceback

from viper.common.out import *
from viper.core.session import __sessions__
from viper.core.plugins import __modules__
from viper.core.project import __project__
from viper.core.ui.commands import Commands
from viper.core.storage import get_sample_path
from viper.core.database import Database
from viper.core.config import Config

cfg = Config()

# For python2 & 3 compat, a bit dirty, but it seems to be the least bad one
try:
    input = raw_input
except NameError:
    pass

def logo():
    print("""         _
        (_)
   _   _ _ ____  _____  ____
  | | | | |  _ \| ___ |/ ___)
   \ V /| | |_| | ____| |
    \_/ |_|  __/|_____)_| v1.3-dev
          |_|
    """)

    db = Database()
    count = db.get_sample_count()

    # Handle the New database format
    try:
        db.find('all', None)
    except:
        print_error("You need to update your viper database. Run 'python update.py -d'")
        sys.exit()

    if __project__.name:
        name = __project__.name
    else:
        name = 'default'

    print(magenta("You have " + bold(count)) +
          magenta(" files in your " + bold(name) +
          magenta(" repository".format(bold(name)))))
    if cfg.autorun.enabled and len(cfg.autorun.commands) == 0:
        print_warning("You have enabled autorun but not set any commands in viper.conf.")

class Console(object):

    def __init__(self):
        # This will keep the main loop active as long as it's set to True.
        self.active = True
        self.cmd = Commands()

    def parse(self, data):
        root = ''
        args = []

        # Split words by white space.
        words = data.split()
        # First word is the root command.
        root = words[0]

        # If there are more words, populate the arguments list.
        if len(words) > 1:
            args = words[1:]

        return (root, args)

    def keywords(self, data):
        # Check if $self is in the user input data.
        if '$self' in data:
            # Check if there is an open session.
            if __sessions__.is_set():
                # If a session is opened, replace $self with the path to
                # the file which is currently being analyzed.
                data = data.replace('$self', __sessions__.current.file.path)
            else:
                print("No session opened")
                return None

        return data

    def stop(self):
        # Stop main loop.
        self.active = False

    def start(self):
        # Logo.
        logo()

        # Setup shell auto-complete.
        def complete(text, state):
            # Try to autocomplete commands.
            cmds = [i for i in self.cmd.commands if i.startswith(text)]
            if state < len(cmds):
                return cmds[state]

            # Try to autocomplete modules.
            mods = [i for i in __modules__ if i.startswith(text)]
            if state < len(mods):
                return mods[state]

            # Then autocomplete paths.
            if text.startswith("~"):
                text = "{0}{1}".format(os.getenv("HOME"), text[1:])
            return (glob.glob(text+'*')+[None])[state]

        # Auto-complete on tabs.
        readline.set_completer_delims(' \t\n;')
        readline.parse_and_bind('tab: complete')
        readline.set_completer(complete)

        # Save commands in history file.
        def save_history(path):
            readline.write_history_file(path)

        # If there is an history file, read from it and load the history
        # so that they can be loaded in the shell.
        # Now we are storing the history file in the local project folder
        # if there is an opened project. Otherwise just store it in the
        # home directory.
        if __project__.path:
            history_path = os.path.join(__project__.path, 'history')
        else:
            history_path = os.path.expanduser('~/.viperhistory')

        if os.path.exists(history_path):
            readline.read_history_file(history_path)

        # Register the save history at program's exit.
        atexit.register(save_history, path=history_path)

        # Main loop.
        while self.active:
            # If there is an open session, we include the path to the opened
            # file in the shell prompt.
            # TODO: perhaps this block should be moved into the session so that
            # the generation of the prompt is done only when the session's
            # status changes.
            prefix = ''
            if __project__.name:
                prefix = bold(cyan(__project__.name)) + ' '

            if __sessions__.is_set():
                stored = ''
                filename = ''
                if __sessions__.current.file:
                    filename = __sessions__.current.file.name
                    if not Database().find(key='sha256', value=__sessions__.current.file.sha256):
                        stored = magenta(' [not stored]', True)
                misp = ''
                if __sessions__.current.misp_event:
                    misp = ' [MISP {}]'.format(__sessions__.current.misp_event.event_id)
                prompt = (prefix + cyan('viper ', True) +
                          white(filename, True) + blue(misp, True) + stored + cyan(' > ', True))
            # Otherwise display the basic prompt.
            else:
                prompt = prefix + cyan('viper > ', True)

            # Wait for input from the user.
            try:
                data = input(prompt).strip()
            except KeyboardInterrupt:
                print("")
            # Terminate on EOF.
            except EOFError:
                self.stop()
                print("")
                continue
            # Parse the input if the user provided any.
            else:
                # If there are recognized keywords, we replace them with
                # their respective value.
                data = self.keywords(data)
                # Skip if the input is empty.
                if not data:
                    continue

                # Check for output redirection
                # If there is a > in the string, we assume the user wants to output to file.
                filename = False
                if '>' in data:
                    data, filename = data.split('>')

                # If the input starts with an exclamation mark, we treat the
                # input as a bash command and execute it.
                # At this point the keywords should be replaced.
                if data.startswith('!'):
                    os.system(data[1:])
                    continue

                # Try to split commands by ; so that you can sequence multiple
                # commands at once.
                # For example:
                # viper > find name *.pdf; open --last 1; pdf id
                # This will automatically search for all PDF files, open the first entry
                # and run the pdf module against it.
                split_commands = data.split(';')
                for split_command in split_commands:
                    split_command = split_command.strip()
                    if not split_command:
                        continue

                    # If it's an internal command, we parse the input and split it
                    # between root command and arguments.
                    root, args = self.parse(split_command)

                    # Check if the command instructs to terminate.
                    if root in ('exit', 'quit'):
                        self.stop()
                        continue

                    try:
                        # If the root command is part of the embedded commands list we
                        # execute it.
                        if root in self.cmd.commands:
                            self.cmd.commands[root]['obj'](*args)
                            print_output(self.cmd.output, filename)
                            del(self.cmd.output[:])
                        # If the root command is part of loaded modules, we initialize
                        # the module and execute it.
                        elif root in __modules__:
                            module = __modules__[root]['obj']()
                            module.set_commandline(args)
                            module.run()

                            print_output(module.output, filename)
                            if cfg.modules.store_output and __sessions__.is_set():
                                Database().add_analysis(__sessions__.current.file.sha256, split_command, module.output)
                            del(module.output[:])
                        else:
                            print("Command not recognized.")
                    except KeyboardInterrupt:
                        pass
                    except Exception as e:
                        print_error("The command {0} raised an exception:".format(bold(root)))
                        traceback.print_exc()
<commit_msg>Error Detecting valid session<commit_after># This file is part of Viper - https://github.com/viper-framework/viper
# See the file 'LICENSE' for copying permission.

import os
import sys
import glob
import atexit
import readline
import traceback

from viper.common.out import *
from viper.core.session import __sessions__
from viper.core.plugins import __modules__
from viper.core.project import __project__
from viper.core.ui.commands import Commands
from viper.core.storage import get_sample_path
from viper.core.database import Database
from viper.core.config import Config

cfg = Config()

# For python2 & 3 compat, a bit dirty, but it seems to be the least bad one
try:
    input = raw_input
except NameError:
    pass

def logo():
    print("""         _
        (_)
   _   _ _ ____  _____  ____
  | | | | |  _ \| ___ |/ ___)
   \ V /| | |_| | ____| |
    \_/ |_|  __/|_____)_| v1.3-dev
          |_|
    """)

    db = Database()
    count = db.get_sample_count()

    # Handle the New database format
    try:
        db.find('all', None)
    except:
        print_error("You need to update your viper database. Run 'python update.py -d'")
        sys.exit()

    if __project__.name:
        name = __project__.name
    else:
        name = 'default'

    print(magenta("You have " + bold(count)) +
          magenta(" files in your " + bold(name) +
          magenta(" repository".format(bold(name)))))
    if cfg.autorun.enabled and len(cfg.autorun.commands) == 0:
        print_warning("You have enabled autorun but not set any commands in viper.conf.")

class Console(object):

    def __init__(self):
        # This will keep the main loop active as long as it's set to True.
        self.active = True
        self.cmd = Commands()

    def parse(self, data):
        root = ''
        args = []

        # Split words by white space.
        words = data.split()
        # First word is the root command.
        root = words[0]

        # If there are more words, populate the arguments list.
        if len(words) > 1:
            args = words[1:]

        return (root, args)

    def keywords(self, data):
        # Check if $self is in the user input data.
        if '$self' in data:
            # Check if there is an open session.
            if __sessions__.is_set():
                # If a session is opened, replace $self with the path to
                # the file which is currently being analyzed.
                data = data.replace('$self', __sessions__.current.file.path)
            else:
                print("No session opened")
                return None

        return data

    def stop(self):
        # Stop main loop.
        self.active = False

    def start(self):
        # Logo.
        logo()

        # Setup shell auto-complete.
        def complete(text, state):
            # Try to autocomplete commands.
            cmds = [i for i in self.cmd.commands if i.startswith(text)]
            if state < len(cmds):
                return cmds[state]

            # Try to autocomplete modules.
            mods = [i for i in __modules__ if i.startswith(text)]
            if state < len(mods):
                return mods[state]

            # Then autocomplete paths.
            if text.startswith("~"):
                text = "{0}{1}".format(os.getenv("HOME"), text[1:])
            return (glob.glob(text+'*')+[None])[state]

        # Auto-complete on tabs.
        readline.set_completer_delims(' \t\n;')
        readline.parse_and_bind('tab: complete')
        readline.set_completer(complete)

        # Save commands in history file.
        def save_history(path):
            readline.write_history_file(path)

        # If there is an history file, read from it and load the history
        # so that they can be loaded in the shell.
        # Now we are storing the history file in the local project folder
        # if there is an opened project. Otherwise just store it in the
        # home directory.
        if __project__.path:
            history_path = os.path.join(__project__.path, 'history')
        else:
            history_path = os.path.expanduser('~/.viperhistory')

        if os.path.exists(history_path):
            readline.read_history_file(history_path)

        # Register the save history at program's exit.
        atexit.register(save_history, path=history_path)

        # Main loop.
        while self.active:
            # If there is an open session, we include the path to the opened
            # file in the shell prompt.
            # TODO: perhaps this block should be moved into the session so that
            # the generation of the prompt is done only when the session's
            # status changes.
            prefix = ''
            if __project__.name:
                prefix = bold(cyan(__project__.name)) + ' '

            if __sessions__.is_set():
                stored = ''
                filename = ''
                if __sessions__.current.file:
                    filename = __sessions__.current.file.name
                    if not Database().find(key='sha256', value=__sessions__.current.file.sha256):
                        stored = magenta(' [not stored]', True)
                misp = ''
                if __sessions__.current.misp_event:
                    misp = ' [MISP {}]'.format(__sessions__.current.misp_event.event_id)
                prompt = (prefix + cyan('viper ', True) +
                          white(filename, True) + blue(misp, True) + stored + cyan(' > ', True))
            # Otherwise display the basic prompt.
            else:
                prompt = prefix + cyan('viper > ', True)

            # Wait for input from the user.
            try:
                data = input(prompt).strip()
            except KeyboardInterrupt:
                print("")
            # Terminate on EOF.
            except EOFError:
                self.stop()
                print("")
                continue
            # Parse the input if the user provided any.
            else:
                # If there are recognized keywords, we replace them with
                # their respective value.
                data = self.keywords(data)
                # Skip if the input is empty.
                if not data:
                    continue

                # Check for output redirection
                # If there is a > in the string, we assume the user wants to output to file.
                filename = False
                if '>' in data:
                    data, filename = data.split('>')

                # If the input starts with an exclamation mark, we treat the
                # input as a bash command and execute it.
                # At this point the keywords should be replaced.
                if data.startswith('!'):
                    os.system(data[1:])
                    continue

                # Try to split commands by ; so that you can sequence multiple
                # commands at once.
                # For example:
                # viper > find name *.pdf; open --last 1; pdf id
                # This will automatically search for all PDF files, open the first entry
                # and run the pdf module against it.
                split_commands = data.split(';')
                for split_command in split_commands:
                    split_command = split_command.strip()
                    if not split_command:
                        continue

                    # If it's an internal command, we parse the input and split it
                    # between root command and arguments.
                    root, args = self.parse(split_command)

                    # Check if the command instructs to terminate.
                    if root in ('exit', 'quit'):
                        self.stop()
                        continue

                    try:
                        # If the root command is part of the embedded commands list we
                        # execute it.
                        if root in self.cmd.commands:
                            self.cmd.commands[root]['obj'](*args)
                            print_output(self.cmd.output, filename)
                            del(self.cmd.output[:])
                        # If the root command is part of loaded modules, we initialize
                        # the module and execute it.
                        elif root in __modules__:
                            module = __modules__[root]['obj']()
                            module.set_commandline(args)
                            module.run()

                            print_output(module.output, filename)
                            if cfg.modules.store_output and __sessions__.current.file:
                                Database().add_analysis(__sessions__.current.file.sha256, split_command, module.output)
                            del(module.output[:])
                        else:
                            print("Command not recognized.")
                    except KeyboardInterrupt:
                        pass
                    except Exception as e:
                        print_error("The command {0} raised an exception:".format(bold(root)))
                        traceback.print_exc()
<|endoftext|>


<commit_before># MIT license
#
# Copyright (C) 2015 by XESS Corporation
#
# 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.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from .kicad import *
<commit_msg>Delete __init__.py<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python
import cgi
import cgitb; cgitb.enable()  # for troubleshooting
import os
from hashlib import sha256
import hmac
from collections import defaultdict

SHARED_SECRET_KEY = 'BESURETODRINKYOUROVALTINE'
def authenticate_with_codepuppy(token, forward_url, verify_url):
  sunetid = os.environ['REMOTE_USER']
  h = hmac.new(SHARED_SECRET_KEY, "{0},{1}".format(token, sunetid), sha256)
  url = "{0}?token={1}&user_system_id={2}&hmac={3}&foward_url={4}"
  url = url.format(verify_url, token, sunetid, h.hexdigest(), forward_url)
  print("Content-Type: text/plain")
  print("")
  print("Location:{0}".format(url))

def main():
  form = cgi.FieldStorage()
  token = form.getvalue("token")
  forward_url = form.getvalue("forward_url")
  verify_url = form.getvalue("verify_url")
  if token == None:
    print("Content-Type: text/plain")
    print("")
    print("unauthorized user plz go")
    return
  authenticate_with_codepuppy(token, forward_url, verify_url)

if __name__ == '__main__':
  main()
<commit_msg>Fixed CGI auth script<commit_after>#!/usr/bin/env python
import cgi
import cgitb; cgitb.enable()  # for troubleshooting
import os
from hashlib import sha256
import hmac
from collections import defaultdict

SHARED_SECRET_KEY = 'BESURETODRINKYOUROVALTINE'
def authenticate_with_codepuppy(token, forward_url, verify_url):
  sunetid = os.environ['REMOTE_USER']
  h = hmac.new(SHARED_SECRET_KEY, "{0},{1}".format(token, sunetid), sha256)
  url = "{0}?token={1}&user_system_id={2}&hmac={3}&forward_url={4}"
  url = url.format(verify_url, token, sunetid, h.hexdigest(), forward_url)
  print("Location:{0}".format(url))
  print("")

def main():
  form = cgi.FieldStorage()
  token = form.getvalue("token")
  forward_url = form.getvalue("forward_url")
  verify_url = form.getvalue("verify_url")
  if token == None:
    print("Content-Type: text/plain")
    print("")
    print("unauthorized user plz go")
    return
  authenticate_with_codepuppy(token, forward_url, verify_url)

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


<commit_before># -*- coding: utf-8 -*-
# This file is part of Viper - https://github.com/viper-framework/viper
# and developed by Koodous Team.
# See the file 'LICENSE' for copying permission.

try:
    import requests

    HAVE_REQUESTS = True
except ImportError:
    HAVE_REQUESTS = False

from io import BytesIO
import hashlib
import tempfile

from viper.common.abstracts import Module
from viper.core.session import __sessions__
from viper.core.config import Config

cfg = Config()

class Koodous(Module):
    cmd = 'koodous'
    description = 'Interact with Koodous'
    authors = ['asanchez@koodous.com']

    def __init__(self):
        super(Koodous, self).__init__()
        self.parser.add_argument('-d', '--download', action='store',
                                 dest='sha256')
        self.parser.add_argument('-u', '--upload', action='store_true',
                                 help='Upload file to Koodous')
        self.parser.add_argument('-c', '--comment', nargs='+', action='store',
                                 dest='comment',
                                 help='Make a comment about sample')
        self.parser.add_argument('-lc', '--load-comments', action='store_true',
                                 help='Show comments about the sample loaded')

    def run(self):
        super(Koodous, self).run()
        if self.args is None:
            return

        if not HAVE_REQUESTS:
            self.log('error', "Missing dependency, install requests (`pip install requests`)")
            return

        if self.args.sha256:
            self._download(self.args.sha256)
            return

        if self.args.upload:
            self._upload()
            return

        if self.args.comment:
            self._comment(' '.join(self.args.comment))
            return

        if self.args.load_comments:
            self._show_comments()
            return

        # If no action, show the help
        self.help()

    def _download(self, sha256):
        """
            Function to download a sample from Koodous
        """
        url = '%s/%s/download' % (cfg.koodous.base_url, sha256)

        headers = {'Authorization': 'Token %s' % cfg.koodous.token}

        response = requests.get(url=url, headers=headers)
        if response.status_code == 200:
            down_url = response.json().get('download_url', None)
            response = requests.get(url=down_url)

            sha256_downloaded = hashlib.sha256(response.content).hexdigest()
            if sha256_downloaded != sha256:
                self.log('error', 'Problem downloading')
            tmp = tempfile.NamedTemporaryFile(delete=False)
            tmp.write(response.content)
            tmp.close()
            return __sessions__.new(tmp.name)

    def _upload(self):
        """
            Function to upload the session file to Koodous
        """

        content_file = __sessions__.current.file.data
        if content_file[:2] != b'PK':
            self.log('info', 'Koodous only accepts APKs, try with VirusTotal.')
            return
        sha256 = hashlib.sha256(content_file).hexdigest()

        url = '%s/%s/get_upload_url' % (cfg.koodous.koodous_url, sha256)
        headers = {"Authorization": "Token %s" % cfg.koodous.koodous_token}

        response = requests.get(url=url, headers=headers)
        if response.status_code == 200:
            upload_url = response.json().get('upload_url', None)
            files = {'file': BytesIO(__sessions__.current.file.data)}
            response = requests.post(url=upload_url, files=files)
            if response == 200:
                self.log("File uploaded correctly.")
                return
        elif response.status_code == 409:
            self.log('info', 'This file already exists in Koodous.')
        else:
            self.log('error', 'Unknown error, sorry!')

    def _comment(self, comment):
        """
            Function to comment an APK in Koodous
        """
        headers = {"Authorization": "Token %s" % cfg.koodous.koodous_token}

        try:
            content_file = __sessions__.current.file.data
            sha256 = hashlib.sha256(content_file).hexdigest()
        except:
            self.log('error', 'You have no file loaded')

        try:
            url = '%s/%s/comments' % (cfg.koodous.koodous_url, sha256)
            data = {'text': comment}
            response = requests.post(url=url, headers=headers, data=data)
            if response.status_code == 201:
                self.log('info', 'Comment made successfully.')
            else:
                self.log('error', 'Some problem saving the comment, please try again.')
        except:
            self.log('error', 'Network problem, please try again.')

    def _show_comments(self):
        """
            Function to view comments of a sample in Koodous
        """
        headers = {"Authorization": "Token %s" % cfg.koodous.koodous_token}

        try:
            content_file = __sessions__.current.file.data
            sha256 = hashlib.sha256(content_file).hexdigest()
        except:
            self.log('error', 'You have no file loaded')
            return
        try:
            url = '%s/%s/comments' % (cfg.koodous.koodous_url, sha256)
            response = requests.get(url=url, headers=headers)
            if response.json().get('count') == 0:
                self.log('info', 'This sample has no comments.')
                return
            for result in response.json().get('results'):
                self.log('info', "[{}]: {}".format(result['author']['username'], result['text']))
        except:
            self.log('error', 'Network problem, please try again.')
<commit_msg>Update Koodus to the new way to access the file content<commit_after># This file is part of Viper - https://github.com/viper-framework/viper
# and developed by Koodous Team.
# See the file 'LICENSE' for copying permission.

try:
    import requests
    HAVE_REQUESTS = True
except ImportError:
    HAVE_REQUESTS = False

from io import BytesIO
import hashlib
import tempfile

from viper.common.abstracts import Module
from viper.core.session import __sessions__
from viper.core.config import Config

cfg = Config()

class Koodous(Module):
    cmd = 'koodous'
    description = 'Interact with Koodous'
    authors = ['asanchez@koodous.com']

    def __init__(self):
        super(Koodous, self).__init__()
        self.parser.add_argument('-d', '--download', action='store',
                                 dest='sha256')
        self.parser.add_argument('-u', '--upload', action='store_true',
                                 help='Upload file to Koodous')
        self.parser.add_argument('-c', '--comment', nargs='+', action='store',
                                 dest='comment',
                                 help='Make a comment about sample')
        self.parser.add_argument('-lc', '--load-comments', action='store_true',
                                 help='Show comments about the sample loaded')

    def run(self):
        super(Koodous, self).run()
        if self.args is None:
            return

        if not HAVE_REQUESTS:
            self.log('error', "Missing dependency, install requests (`pip install requests`)")
            return

        if self.args.sha256:
            self._download(self.args.sha256)
            return

        if self.args.upload:
            self._upload()
            return

        if self.args.comment:
            self._comment(' '.join(self.args.comment))
            return

        if self.args.load_comments:
            self._show_comments()
            return

        # If no action, show the help
        self.help()

    def _download(self, sha256):
        """
            Function to download a sample from Koodous
        """
        url = '%s/%s/download' % (cfg.koodous.base_url, sha256)

        headers = {'Authorization': 'Token %s' % cfg.koodous.token}

        response = requests.get(url=url, headers=headers)
        if response.status_code == 200:
            down_url = response.json().get('download_url', None)
            response = requests.get(url=down_url)

            sha256_downloaded = hashlib.sha256(response.content).hexdigest()
            if sha256_downloaded != sha256:
                self.log('error', 'Problem downloading')
            tmp = tempfile.NamedTemporaryFile(delete=False)
            tmp.write(response.content)
            tmp.close()
            return __sessions__.new(tmp.name)

    def _upload(self):
        """
            Function to upload the session file to Koodous
        """

        content_file = __sessions__.current.file.data
        if content_file[:2] != b'PK':
            self.log('info', 'Koodous only accepts APKs, try with VirusTotal.')
            return
        sha256 = hashlib.sha256(content_file).hexdigest()

        url = '%s/%s/get_upload_url' % (cfg.koodous.koodous_url, sha256)
        headers = {"Authorization": "Token %s" % cfg.koodous.koodous_token}

        response = requests.get(url=url, headers=headers)
        if response.status_code == 200:
            upload_url = response.json().get('upload_url', None)
            files = {'file': ByteIO(__sessions__.current.file.data)}
            response = requests.post(url=upload_url, files=files)
            if response == 200:
                self.log("File uploaded correctly.")
                return
        elif response.status_code == 409:
            self.log('info', 'This file already exists in Koodous.')
        else:
            self.log('error', 'Unknown error, sorry!')

    def _comment(self, comment):
        """
            Function to comment an APK in Koodous
        """
        headers = {"Authorization": "Token %s" % cfg.koodous.koodous_token}

        try:
            content_file = __sessions__.current.file.data
            sha256 = hashlib.sha256(content_file).hexdigest()
        except:
            self.log('error', 'You have no file loaded')

        try:
            url = '%s/%s/comments' % (cfg.koodous.koodous_url, sha256)
            data = {'text': comment}
            response = requests.post(url=url, headers=headers, data=data)
            if response.status_code == 201:
                self.log('info', 'Comment made successfully.')
            else:
                self.log('error', 'Some problem saving the comment, please try again.')
        except:
            self.log('error', 'Network problem, please try again.')

    def _show_comments(self):
        """
            Function to view comments of a sample in Koodous
        """
        headers = {"Authorization": "Token %s" % cfg.koodous.koodous_token}

        try:
            content_file = __sessions__.current.file.data
            sha256 = hashlib.sha256(content_file).hexdigest()
        except:
            self.log('error', 'You have no file loaded')
            return
        try:
            url = '%s/%s/comments' % (cfg.koodous.koodous_url, sha256)
            response = requests.get(url=url, headers=headers)
            if response.json().get('count') == 0:
                self.log('info', 'This sample has no comments.')
                return
            for result in response.json().get('results'):
                self.log('info', "[{}]: {}".format(result['author']['username'], result['text']))
        except:
            self.log('error', 'Network problem, please try again.')
<|endoftext|>


<commit_before>"""Compatibility fixes for older version of python, numpy and scipy"""
# Authors: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
#          Gael Varoquaux <gael.varoquaux@normalesup.org>
#          Fabian Pedregosa <fpedregosa@acm.org>
#          Lars Buitinck <L.J.Buitinck@uva.nl>
# License: BSD

import collections
import numpy as np
from operator import itemgetter

try:
    Counter = collections.Counter
except AttributeError:
# Partial replacement for Python 2.7 collections.Counter
    class Counter(collections.defaultdict):
        def __init__(self, iterable=(), **kwargs):
            super(Counter, self).__init__(int, **kwargs)
            self.update(iterable)

        def most_common(self):
            return sorted(self.iteritems(), key=itemgetter(1), reverse=True)

        def update(self, other):
            """Adds counts for elements in other"""
            if isinstance(other, self.__class__):
                for x, n in other.iteritems():
                    self[x] += n
            else:
                for x in other:
                    self[x] += 1

def _unique(ar, return_index=False, return_inverse=False):
    """A replacement for the np.unique that appeared in numpy 1.4.

    While np.unique existed long before, keyword return_inverse was
    only added in 1.4.
    """
    try:
        ar = ar.flatten()
    except AttributeError:
        if not return_inverse and not return_index:
            items = sorted(set(ar))
            return np.asarray(items)
        else:
            ar = np.asarray(ar).flatten()

    if ar.size == 0:
        if return_inverse and return_index:
            return ar, np.empty(0, np.bool), np.empty(0, np.bool)
        elif return_inverse or return_index:
            return ar, np.empty(0, np.bool)
        else:
            return ar

    if return_inverse or return_index:
        perm = ar.argsort()
        aux = ar[perm]
        flag = np.concatenate(([True], aux[1:] != aux[:-1]))
        if return_inverse:
            iflag = np.cumsum(flag) - 1
            iperm = perm.argsort()
            if return_index:
                return aux[flag], perm[flag], iflag[iperm]
            else:
                return aux[flag], iflag[iperm]
        else:
            return aux[flag], perm[flag]

    else:
        ar.sort()
        flag = np.concatenate(([True], ar[1:] != ar[:-1]))
        return ar[flag]

np_version = np.__version__.split('.')
if int(np_version[0]) < 2 and int(np_version[1]) < 5:
    unique = _unique
else:
    unique = np.unique

def _copysign(x1, x2):
    """Slow replacement for np.copysign, which was introduced in numpy 1.4"""
    return np.abs(x1) * np.sign(x2)

if not hasattr(np, 'copysign'):
    copysign = _copysign
else:
    copysign = np.copysign

def _in1d(ar1, ar2, assume_unique=False):
    """Replacement for in1d that is provided for numpy >= 1.4"""
    if not assume_unique:
        ar1, rev_idx = unique(ar1, return_inverse=True)
        ar2 = np.unique(ar2)
    ar = np.concatenate((ar1, ar2))
    # We need this to be a stable sort, so always use 'mergesort'
    # here. The values from the first array should always come before
    # the values from the second array.
    order = ar.argsort(kind='mergesort')
    sar = ar[order]
    equal_adj = (sar[1:] == sar[:-1])
    flag = np.concatenate((equal_adj, [False]))
    indx = order.argsort(kind='mergesort')[:len(ar1)]

    if assume_unique:
        return flag[indx]
    else:
        return flag[indx][rev_idx]

if not hasattr(np, 'in1d'):
    in1d = _in1d
else:
    in1d = np.in1d

def qr_economic(A, **kwargs):
    """Compat function for the QR-decomposition in economic mode

    Scipy 0.9 changed the keyword econ=True to mode='economic'
    """
    import scipy.linalg
    # trick: triangular solve has introduced in 0.9
    if hasattr(scipy.linalg, 'solve_triangular'):
        return scipy.linalg.qr(A, mode='economic', **kwargs)
    else:
        return scipy.linalg.qr(A, econ=True, **kwargs)

def arpack_eigsh(A, **kwargs):
    """Compat function for sparse symmetric eigen vectors decomposition

    Scipy 0.9 renamed eigen_symmetric to eigsh in
    scipy.sparse.linalg.eigen.arpack
    """
    from scipy.sparse.linalg.eigen import arpack
    if hasattr(arpack, 'eigsh'):
        return arpack.eigsh(A, **kwargs)
    else:
        return arpack.eigen_symmetric(A, **kwargs)

def savemat(file_name, mdict, oned_as="column", **kwargs):
    """MATLAB-format output routine that is compatible with SciPy 0.7's.

    0.7.2 (or .1?) added the oned_as keyword arg with 'column' as the default
    value. It issues a warning if this is not provided, stating that "This will
    change to 'row' in future versions."
    """
    import scipy.io
    try:
        return scipy.io.savemat(file_name, mdict, oned_as=oned_as, **kwargs)
    except TypeError:
        return scipy.io.savemat(file_name, mdict, **kwargs)

try:
    from numpy import count_nonzero
except ImportError:
    def count_nonzero(X):
        return len(np.flatnonzero(X))
<commit_msg>COSMIT: utils.fixes: document versions<commit_after>"""Compatibility fixes for older version of python, numpy and scipy

If you add content to this file, please give the version of the package
at which the fixe is no longer needed.
"""
# Authors: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
#          Gael Varoquaux <gael.varoquaux@normalesup.org>
#          Fabian Pedregosa <fpedregosa@acm.org>
#          Lars Buitinck <L.J.Buitinck@uva.nl>
# License: BSD

import collections
import numpy as np
from operator import itemgetter

# Python < 2.6 does not have product
try:
    Counter = collections.Counter
except AttributeError:
# Partial replacement for Python 2.7 collections.Counter
    class Counter(collections.defaultdict):
        def __init__(self, iterable=(), **kwargs):
            super(Counter, self).__init__(int, **kwargs)
            self.update(iterable)

        def most_common(self):
            return sorted(self.iteritems(), key=itemgetter(1), reverse=True)

        def update(self, other):
            """Adds counts for elements in other"""
            if isinstance(other, self.__class__):
                for x, n in other.iteritems():
                    self[x] += n
            else:
                for x in other:
                    self[x] += 1

def _unique(ar, return_index=False, return_inverse=False):
    """A replacement for the np.unique that appeared in numpy 1.4.

    While np.unique existed long before, keyword return_inverse was
    only added in 1.4.
    """
    try:
        ar = ar.flatten()
    except AttributeError:
        if not return_inverse and not return_index:
            items = sorted(set(ar))
            return np.asarray(items)
        else:
            ar = np.asarray(ar).flatten()

    if ar.size == 0:
        if return_inverse and return_index:
            return ar, np.empty(0, np.bool), np.empty(0, np.bool)
        elif return_inverse or return_index:
            return ar, np.empty(0, np.bool)
        else:
            return ar

    if return_inverse or return_index:
        perm = ar.argsort()
        aux = ar[perm]
        flag = np.concatenate(([True], aux[1:] != aux[:-1]))
        if return_inverse:
            iflag = np.cumsum(flag) - 1
            iperm = perm.argsort()
            if return_index:
                return aux[flag], perm[flag], iflag[iperm]
            else:
                return aux[flag], iflag[iperm]
        else:
            return aux[flag], perm[flag]

    else:
        ar.sort()
        flag = np.concatenate(([True], ar[1:] != ar[:-1]))
        return ar[flag]

np_version = np.__version__.split('.')
if int(np_version[0]) < 2 and int(np_version[1]) < 5:
    unique = _unique
else:
    unique = np.unique

def _copysign(x1, x2):
    """Slow replacement for np.copysign, which was introduced in numpy 1.4"""
    return np.abs(x1) * np.sign(x2)

if not hasattr(np, 'copysign'):
    copysign = _copysign
else:
    copysign = np.copysign

def _in1d(ar1, ar2, assume_unique=False):
    """Replacement for in1d that is provided for numpy >= 1.4"""
    if not assume_unique:
        ar1, rev_idx = unique(ar1, return_inverse=True)
        ar2 = np.unique(ar2)
    ar = np.concatenate((ar1, ar2))
    # We need this to be a stable sort, so always use 'mergesort'
    # here. The values from the first array should always come before
    # the values from the second array.
    order = ar.argsort(kind='mergesort')
    sar = ar[order]
    equal_adj = (sar[1:] == sar[:-1])
    flag = np.concatenate((equal_adj, [False]))
    indx = order.argsort(kind='mergesort')[:len(ar1)]

    if assume_unique:
        return flag[indx]
    else:
        return flag[indx][rev_idx]

if not hasattr(np, 'in1d'):
    in1d = _in1d
else:
    in1d = np.in1d

def qr_economic(A, **kwargs):
    """Compat function for the QR-decomposition in economic mode

    Scipy 0.9 changed the keyword econ=True to mode='economic'
    """
    import scipy.linalg
    # trick: triangular solve has introduced in 0.9
    if hasattr(scipy.linalg, 'solve_triangular'):
        return scipy.linalg.qr(A, mode='economic', **kwargs)
    else:
        return scipy.linalg.qr(A, econ=True, **kwargs)

def arpack_eigsh(A, **kwargs):
    """Compat function for sparse symmetric eigen vectors decomposition

    Scipy 0.9 renamed eigen_symmetric to eigsh in
    scipy.sparse.linalg.eigen.arpack
    """
    from scipy.sparse.linalg.eigen import arpack
    if hasattr(arpack, 'eigsh'):
        return arpack.eigsh(A, **kwargs)
    else:
        return arpack.eigen_symmetric(A, **kwargs)

def savemat(file_name, mdict, oned_as="column", **kwargs):
    """MATLAB-format output routine that is compatible with SciPy 0.7's.

    0.7.2 (or .1?) added the oned_as keyword arg with 'column' as the default
    value. It issues a warning if this is not provided, stating that "This will
    change to 'row' in future versions."
    """
    import scipy.io
    try:
        return scipy.io.savemat(file_name, mdict, oned_as=oned_as, **kwargs)
    except TypeError:
        return scipy.io.savemat(file_name, mdict, **kwargs)

try:
    from numpy import count_nonzero
except ImportError:
    def count_nonzero(X):
        return len(np.flatnonzero(X))
<|endoftext|>


<commit_before>import numpy as np

from sklearn.base import clone
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.utils import check_random_state
from sklearn.externals.joblib import Parallel, delayed

def _parallel_fit(regressor, X, y):
    return regressor.fit(X, y)

class GradientBoostingQuantileRegressor(BaseEstimator, RegressorMixin):
    """Predict several quantiles with one estimator.

    This is a wrapper around `GradientBoostingRegressor`'s quantile
    regression that allows you to predict several `quantiles` in
    one go.

    Parameters
    ----------
    * `quantiles` [array-like]:
        Quantiles to predict. By default the 16, 50 and 84%
        quantiles are predicted.

    * `base_estimator` [GradientBoostingRegressor instance or None (default)]:
        Quantile regressor used to make predictions. Only instances
        of `GradientBoostingRegressor` are supported. Use this to change
        the hyper-parameters of the estimator.

    * `n_jobs` [int, default=1]:
        The number of jobs to run in parallel for `fit`.
        If -1, then the number of jobs is set to the number of cores.

    * `random_state` [int, RandomState instance, or None (default)]:
        Set random state to something other than None for reproducible
        results.
    """

    def __init__(self, quantiles=[0.16, 0.5, 0.84], base_estimator=None,
                 n_jobs=1, random_state=None):
        self.quantiles = quantiles
        self.random_state = random_state
        self.base_estimator = base_estimator
        self.n_jobs = n_jobs

    def fit(self, X, y):
        """Fit one regressor for each quantile.

        Parameters
        ----------
        * `X` [array-like, shape=(n_samples, n_features):
            Training vectors, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        * `y` [array-like, shape=(n_samples,)]:
            Target values (real numbers in regression)
        """
        rng = check_random_state(self.random_state)

        if self.base_estimator is None:
            base_estimator = GradientBoostingRegressor(loss='quantile')
        else:
            base_estimator = self.base_estimator

            if not isinstance(base_estimator, GradientBoostingRegressor):
                raise ValueError('base_estimator has to be of type'
                                 ' GradientBoostingRegressor.')

            if not base_estimator.loss == 'quantile':
                raise ValueError('base_estimator has to use quantile'
                                 ' loss not %s' % base_estimator.loss)

        regressors = []
        for q in self.quantiles:
            regressor = clone(base_estimator)
            regressor.set_params(alpha=q, random_state=rng)

            regressors.append(regressor)

        self.regressors_ = Parallel(n_jobs=self.n_jobs, backend='threading')(
            delayed(_parallel_fit)(regressor, X, y)
            for regressor in regressors)

        return self

    def predict(self, X, return_std=False):
        """Predict.

        Predict `X` at every quantile if `return_std` is set to False.
        If `return_std` is set to True, then return the mean
        and the predicted standard deviation, which is approximated as
        the (0.84th quantile - 0.16th quantile) divided by 2.0

        Parameters
        ----------
        * `X` [array-like, shape=(n_samples, n_features):
            where `n_samples` is the number of samples
            and `n_features` is the number of features.
        """
        predicted_quantiles = np.asarray(
            [rgr.predict(X) for rgr in self.regressors_])
        if not return_std:
            return predicted_quantiles.T
        else:
            std_quantiles = [0.16, 0.5, 0.84]
            is_present_mask = np.in1d(std_quantiles, self.quantiles)
            if not np.all(is_present_mask):
                raise ValueError(
                    "return_std works only if the quantiles during "
                    "instantiation include 0.16, 0.5 and 0.84")
            low = self.regressors_[self.quantiles.index(0.16)].predict(X)
            high = self.regressors_[self.quantiles.index(0.84)].predict(X)
            mean = self.regressors_[self.quantiles.index(0.5)].predict(X)
            return mean, ((high - low) / 2.0)
<commit_msg>Set the same random_state instance for parallel fitting and quantiles sorting<commit_after>import numpy as np

from sklearn.base import clone
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.utils import check_random_state
from sklearn.externals.joblib import Parallel, delayed

def _parallel_fit(regressor, X, y):
    return regressor.fit(X, y)

class GradientBoostingQuantileRegressor(BaseEstimator, RegressorMixin):
    """Predict several quantiles with one estimator.

    This is a wrapper around `GradientBoostingRegressor`'s quantile
    regression that allows you to predict several `quantiles` in
    one go.

    Parameters
    ----------
    * `quantiles` [array-like]:
        Quantiles to predict. By default the 16, 50 and 84%
        quantiles are predicted.

    * `base_estimator` [GradientBoostingRegressor instance or None (default)]:
        Quantile regressor used to make predictions. Only instances
        of `GradientBoostingRegressor` are supported. Use this to change
        the hyper-parameters of the estimator.

    * `n_jobs` [int, default=1]:
        The number of jobs to run in parallel for `fit`.
        If -1, then the number of jobs is set to the number of cores.

    * `random_state` [int, RandomState instance, or None (default)]:
        Set random state to something other than None for reproducible
        results.
    """

    def __init__(self, quantiles=[0.16, 0.5, 0.84], base_estimator=None,
                 n_jobs=1, random_state=None):
        self.quantiles = quantiles
        self.random_state = random_state
        self.base_estimator = base_estimator
        self.n_jobs = n_jobs

    def fit(self, X, y):
        """Fit one regressor for each quantile.

        Parameters
        ----------
        * `X` [array-like, shape=(n_samples, n_features):
            Training vectors, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        * `y` [array-like, shape=(n_samples,)]:
            Target values (real numbers in regression)
        """
        rng = check_random_state(self.random_state)

        if self.base_estimator is None:
            base_estimator = GradientBoostingRegressor(loss='quantile')
        else:
            base_estimator = self.base_estimator

            if not isinstance(base_estimator, GradientBoostingRegressor):
                raise ValueError('base_estimator has to be of type'
                                 ' GradientBoostingRegressor.')

            if not base_estimator.loss == 'quantile':
                raise ValueError('base_estimator has to use quantile'
                                 ' loss not %s' % base_estimator.loss)

        # The predictions for different quantiles should be sorted.
        # Therefore each of the regressors need the same seed.
        base_estimator.set_params(random_state=rng)
        regressors = []
        for q in self.quantiles:
            regressor = clone(base_estimator)
            regressor.set_params(alpha=q)

            regressors.append(regressor)

        self.regressors_ = Parallel(n_jobs=self.n_jobs, backend='threading')(
            delayed(_parallel_fit)(regressor, X, y)
            for regressor in regressors)

        return self

    def predict(self, X, return_std=False):
        """Predict.

        Predict `X` at every quantile if `return_std` is set to False.
        If `return_std` is set to True, then return the mean
        and the predicted standard deviation, which is approximated as
        the (0.84th quantile - 0.16th quantile) divided by 2.0

        Parameters
        ----------
        * `X` [array-like, shape=(n_samples, n_features):
            where `n_samples` is the number of samples
            and `n_features` is the number of features.
        """
        predicted_quantiles = np.asarray(
            [rgr.predict(X) for rgr in self.regressors_])
        if not return_std:
            return predicted_quantiles.T
        else:
            std_quantiles = [0.16, 0.5, 0.84]
            is_present_mask = np.in1d(std_quantiles, self.quantiles)
            if not np.all(is_present_mask):
                raise ValueError(
                    "return_std works only if the quantiles during "
                    "instantiation include 0.16, 0.5 and 0.84")
            low = self.regressors_[self.quantiles.index(0.16)].predict(X)
            high = self.regressors_[self.quantiles.index(0.84)].predict(X)
            mean = self.regressors_[self.quantiles.index(0.5)].predict(X)
            return mean, ((high - low) / 2.0)
<|endoftext|>