<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

import logging
import os
import threading

from funcy import cached_property, wrap_prop

from dvc.config import Config
from dvc.exceptions import DvcException
from dvc.exceptions import ETagMismatchError
from dvc.path_info import CloudURLInfo
from dvc.progress import Tqdm
from dvc.remote.base import RemoteBASE
from dvc.scheme import Schemes

logger = logging.getLogger(__name__)

class RemoteS3(RemoteBASE):
    scheme = Schemes.S3
    path_cls = CloudURLInfo
    REQUIRES = {"boto3": "boto3"}
    PARAM_CHECKSUM = "etag"

    def __init__(self, repo, config):
        super(RemoteS3, self).__init__(repo, config)

        url = config.get(Config.SECTION_REMOTE_URL, "s3://")
        self.path_info = self.path_cls(url)

        self.region = config.get(Config.SECTION_AWS_REGION)

        self.profile = config.get(Config.SECTION_AWS_PROFILE)

        self.endpoint_url = config.get(Config.SECTION_AWS_ENDPOINT_URL)

        if config.get(Config.SECTION_AWS_LIST_OBJECTS):
            self.list_objects_api = "list_objects"
        else:
            self.list_objects_api = "list_objects_v2"

        self.use_ssl = config.get(Config.SECTION_AWS_USE_SSL, True)

        self.extra_args = {}

        self.sse = config.get(Config.SECTION_AWS_SSE, "")
        if self.sse:
            self.extra_args["ServerSideEncryption"] = self.sse

        self.acl = config.get(Config.SECTION_AWS_ACL, "")
        if self.acl:
            self.extra_args["ACL"] = self.acl

        self._append_aws_grants_to_extra_args(config)

        shared_creds = config.get(Config.SECTION_AWS_CREDENTIALPATH)
        if shared_creds:
            os.environ.setdefault("AWS_SHARED_CREDENTIALS_FILE", shared_creds)

    @wrap_prop(threading.Lock())
    @cached_property
    def s3(self):
        import boto3

        session = boto3.session.Session(
            profile_name=self.profile, region_name=self.region
        )

        return session.client(
            "s3", endpoint_url=self.endpoint_url, use_ssl=self.use_ssl
        )

    @classmethod
    def get_etag(cls, s3, bucket, path):
        obj = cls.get_head_object(s3, bucket, path)

        return obj["ETag"].strip('"')

    def get_file_checksum(self, path_info):
        return self.get_etag(self.s3, path_info.bucket, path_info.path)

    @staticmethod
    def get_head_object(s3, bucket, path, *args, **kwargs):

        try:
            obj = s3.head_object(Bucket=bucket, Key=path, *args, **kwargs)
        except Exception as exc:
            raise DvcException(
                "s3://{}/{} does not exist".format(bucket, path), exc
            )
        return obj

    @classmethod
    def _copy_multipart(
        cls, s3, from_info, to_info, size, n_parts, extra_args
    ):
        mpu = s3.create_multipart_upload(
            Bucket=to_info.bucket, Key=to_info.path, **extra_args
        )
        mpu_id = mpu["UploadId"]

        parts = []
        byte_position = 0
        for i in range(1, n_parts + 1):
            obj = cls.get_head_object(
                s3, from_info.bucket, from_info.path, PartNumber=i
            )
            part_size = obj["ContentLength"]
            lastbyte = byte_position + part_size - 1
            if lastbyte > size:
                lastbyte = size - 1

            srange = "bytes={}-{}".format(byte_position, lastbyte)

            part = s3.upload_part_copy(
                Bucket=to_info.bucket,
                Key=to_info.path,
                PartNumber=i,
                UploadId=mpu_id,
                CopySourceRange=srange,
                CopySource={"Bucket": from_info.bucket, "Key": from_info.path},
            )
            parts.append(
                {"PartNumber": i, "ETag": part["CopyPartResult"]["ETag"]}
            )
            byte_position += part_size

        assert n_parts == len(parts)

        s3.complete_multipart_upload(
            Bucket=to_info.bucket,
            Key=to_info.path,
            UploadId=mpu_id,
            MultipartUpload={"Parts": parts},
        )

    @classmethod
    def _copy(cls, s3, from_info, to_info, extra_args):
        # NOTE: object's etag depends on the way it was uploaded to s3 or the
        # way it was copied within the s3. More specifically, it depends on
        # the chunk size that was used to transfer it, which would affect
        # whether an object would be uploaded as a single part or as a
        # multipart.
        #
        # If an object's etag looks like '8978c98bb5a48c2fb5f2c4c905768afa',
        # then it was transferred as a single part, which means that the chunk
        # size used to transfer it was greater or equal to the ContentLength
        # of that object. So to preserve that tag over the next transfer, we
        # could use any value >= ContentLength.
        #
        # If an object's etag looks like '50d67013a5e1a4070bef1fc8eea4d5f9-13',
        # then it was transferred as a multipart, which means that the chunk
        # size used to transfer it was less than ContentLength of that object.
        # Unfortunately, in general, it doesn't mean that the chunk size was
        # the same throughout the transfer, so it means that in order to
        # preserve etag, we need to transfer each part separately, so the
        # object is transfered in the same chunks as it was originally.

        obj = cls.get_head_object(s3, from_info.bucket, from_info.path)
        etag = obj["ETag"].strip('"')
        size = obj["ContentLength"]

        _, _, parts_suffix = etag.partition("-")
        if parts_suffix:
            n_parts = int(parts_suffix)
            cls._copy_multipart(
                s3, from_info, to_info, size, n_parts, extra_args=extra_args
            )
        else:
            source = {"Bucket": from_info.bucket, "Key": from_info.path}
            s3.copy(source, to_info.bucket, to_info.path, ExtraArgs=extra_args)

        cached_etag = cls.get_etag(s3, to_info.bucket, to_info.path)
        if etag != cached_etag:
            raise ETagMismatchError(etag, cached_etag)

    def copy(self, from_info, to_info):
        self._copy(self.s3, from_info, to_info, self.extra_args)

    def remove(self, path_info):
        if path_info.scheme != "s3":
            raise NotImplementedError

        logger.debug("Removing {}".format(path_info))
        self.s3.delete_object(Bucket=path_info.bucket, Key=path_info.path)

    def _list_objects(self, path_info, max_items=None):
        """ Read config for list object api, paginate through list objects."""
        kwargs = {
            "Bucket": path_info.bucket,
            "Prefix": path_info.path,
            "PaginationConfig": {"MaxItems": max_items},
        }
        paginator = self.s3.get_paginator(self.list_objects_api)
        for page in paginator.paginate(**kwargs):
            contents = page.get("Contents", None)
            if not contents:
                continue
            for item in contents:
                yield item

    def _list_paths(self, path_info, max_items=None):
        return (
            item["Key"] for item in self._list_objects(path_info, max_items)
        )

    def list_cache_paths(self):
        return self._list_paths(self.path_info)

    def isfile(self, path_info):
        from botocore.exceptions import ClientError

        if path_info.path.endswith("/"):
            return False

        try:
            self.s3.head_object(Bucket=path_info.bucket, Key=path_info.path)
        except ClientError as exc:
            if exc.response["Error"]["Code"] != "404":
                raise
            return False

        return True

    def exists(self, path_info):
        """Check if the blob exists. If it does not exist,
        it could be a part of a directory path.

        eg: if `data/file.txt` exists, check for `data` should return True
        """
        return self.isfile(path_info) or self.isdir(path_info)

    def makedirs(self, path_info):
        # We need to support creating empty directories, which means
        # creating an object with an empty body and a trailing slash `/`.
        #
        # We are not creating directory objects for every parent prefix,
        # as it is not required.
        dir_path = path_info / ""
        self.s3.put_object(Bucket=path_info.bucket, Key=dir_path.path, Body="")

    def isdir(self, path_info):
        # S3 doesn't have a concept for directories.
        #
        # Using `head_object` with a path pointing to a directory
        # will throw a 404 error.
        #
        # A reliable way to know if a given path is a directory is by
        # checking if there are more files sharing the same prefix
        # with a `list_objects` call.
        #
        # We need to make sure that the path ends with a forward slash,
        # since we can end with false-positives like the following example:
        #
        #       bucket
        #       └── data
        #          ├── alice
        #          └── alpha
        #
        # Using `data/al` as prefix will return `[data/alice, data/alpha]`,
        # While `data/al/` will return nothing.
        #
        dir_path = path_info / ""
        return bool(list(self._list_paths(dir_path, max_items=1)))

    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
        total = os.path.getsize(from_file)
        with Tqdm(
            disable=no_progress_bar, total=total, bytes=True, desc=name
        ) as pbar:
            self.s3.upload_file(
                from_file,
                to_info.bucket,
                to_info.path,
                Callback=pbar.update,
                ExtraArgs=self.extra_args,
            )

    def _download(self, from_info, to_file, name=None, no_progress_bar=False):
        if no_progress_bar:
            total = None
        else:
            total = self.s3.head_object(
                Bucket=from_info.bucket, Key=from_info.path
            )["ContentLength"]
        with Tqdm(
            disable=no_progress_bar, total=total, bytes=True, desc=name
        ) as pbar:
            self.s3.download_file(
                from_info.bucket, from_info.path, to_file, Callback=pbar.update
            )

    def _generate_download_url(self, path_info, expires=3600):
        params = {"Bucket": path_info.bucket, "Key": path_info.path}
        return self.s3.generate_presigned_url(
            ClientMethod="get_object", Params=params, ExpiresIn=int(expires)
        )

    def walk_files(self, path_info, max_items=None):
        for fname in self._list_paths(path_info / "", max_items):
            if fname.endswith("/"):
                continue

            yield path_info.replace(path=fname)

    def _append_aws_grants_to_extra_args(self, config):
        # Keys for extra_args can be one of the following list:
        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html#boto3.s3.transfer.S3Transfer.ALLOWED_UPLOAD_ARGS
        """
          ALLOWED_UPLOAD_ARGS = [
            'ACL', 'CacheControl', 'ContentDisposition', 'ContentEncoding',
            'ContentLanguage', 'ContentType', 'Expires', 'GrantFullControl',
            'GrantRead', 'GrantReadACP', 'GrantWriteACP', 'Metadata',
            'RequestPayer', 'ServerSideEncryption', 'StorageClass',
            'SSECustomerAlgorithm', 'SSECustomerKey', 'SSECustomerKeyMD5',
            'SSEKMSKeyId', 'WebsiteRedirectLocation'
          ]
        """
        overwriting_acl_grantees = False

        self.grant_full_control = config.get(
            Config.SECTION_AWS_GRANT_FULL_CONTROL, ""
        )
        if self.grant_full_control:
            self.extra_args["GrantFullControl"] = self.grant_full_control
            overwriting_acl_grantees = True

        self.grant_read = config.get(Config.SECTION_AWS_GRANT_READ, "")
        if self.grant_read:
            self.extra_args["GrantRead"] = self.grant_read
            overwriting_acl_grantees = True

        self.grant_read_acp = config.get(Config.SECTION_AWS_GRANT_READ_ACP, "")
        if self.grant_read_acp:
            self.extra_args["GrantReadACP"] = self.grant_read_acp
            overwriting_acl_grantees = True

        self.grant_write_acp = config.get(
            Config.SECTION_AWS_GRANT_WRITE_ACP, ""
        )
        if self.grant_write_acp:
            self.extra_args["GrantWriteACP"] = self.grant_write_acp
            overwriting_acl_grantees = True

        if overwriting_acl_grantees:
            overwriting_fields = [
                Config.SECTION_AWS_GRANT_FULL_CONTROL,
                Config.SECTION_AWS_GRANT_READ,
                Config.SECTION_AWS_GRANT_READ_ACP,
                Config.SECTION_AWS_GRANT_WRITE_ACP,
            ]
            warning_message = [
                f"Using one of the [{', '.join(overwriting_fields)}] will",
                "overwrite default ACL Grantees list. Default ACL Grantees",
                "list include Owner with FULL_CONTROL permissions.",
            ]
            logger.warning(" ".join(warning_message))
<commit_msg>fixup! remote: Extend S3 remote support for custom ACL<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

import logging
import os
import threading

from funcy import cached_property, wrap_prop

from dvc.config import Config
from dvc.exceptions import DvcException
from dvc.exceptions import ETagMismatchError
from dvc.path_info import CloudURLInfo
from dvc.progress import Tqdm
from dvc.remote.base import RemoteBASE
from dvc.scheme import Schemes

logger = logging.getLogger(__name__)

class RemoteS3(RemoteBASE):
    scheme = Schemes.S3
    path_cls = CloudURLInfo
    REQUIRES = {"boto3": "boto3"}
    PARAM_CHECKSUM = "etag"

    def __init__(self, repo, config):
        super(RemoteS3, self).__init__(repo, config)

        url = config.get(Config.SECTION_REMOTE_URL, "s3://")
        self.path_info = self.path_cls(url)

        self.region = config.get(Config.SECTION_AWS_REGION)

        self.profile = config.get(Config.SECTION_AWS_PROFILE)

        self.endpoint_url = config.get(Config.SECTION_AWS_ENDPOINT_URL)

        if config.get(Config.SECTION_AWS_LIST_OBJECTS):
            self.list_objects_api = "list_objects"
        else:
            self.list_objects_api = "list_objects_v2"

        self.use_ssl = config.get(Config.SECTION_AWS_USE_SSL, True)

        self.extra_args = {}

        self.sse = config.get(Config.SECTION_AWS_SSE, "")
        if self.sse:
            self.extra_args["ServerSideEncryption"] = self.sse

        self.acl = config.get(Config.SECTION_AWS_ACL, "")
        if self.acl:
            self.extra_args["ACL"] = self.acl

        self._append_aws_grants_to_extra_args(config)

        shared_creds = config.get(Config.SECTION_AWS_CREDENTIALPATH)
        if shared_creds:
            os.environ.setdefault("AWS_SHARED_CREDENTIALS_FILE", shared_creds)

    @wrap_prop(threading.Lock())
    @cached_property
    def s3(self):
        import boto3

        session = boto3.session.Session(
            profile_name=self.profile, region_name=self.region
        )

        return session.client(
            "s3", endpoint_url=self.endpoint_url, use_ssl=self.use_ssl
        )

    @classmethod
    def get_etag(cls, s3, bucket, path):
        obj = cls.get_head_object(s3, bucket, path)

        return obj["ETag"].strip('"')

    def get_file_checksum(self, path_info):
        return self.get_etag(self.s3, path_info.bucket, path_info.path)

    @staticmethod
    def get_head_object(s3, bucket, path, *args, **kwargs):

        try:
            obj = s3.head_object(Bucket=bucket, Key=path, *args, **kwargs)
        except Exception as exc:
            raise DvcException(
                "s3://{}/{} does not exist".format(bucket, path), exc
            )
        return obj

    @classmethod
    def _copy_multipart(
        cls, s3, from_info, to_info, size, n_parts, extra_args
    ):
        mpu = s3.create_multipart_upload(
            Bucket=to_info.bucket, Key=to_info.path, **extra_args
        )
        mpu_id = mpu["UploadId"]

        parts = []
        byte_position = 0
        for i in range(1, n_parts + 1):
            obj = cls.get_head_object(
                s3, from_info.bucket, from_info.path, PartNumber=i
            )
            part_size = obj["ContentLength"]
            lastbyte = byte_position + part_size - 1
            if lastbyte > size:
                lastbyte = size - 1

            srange = "bytes={}-{}".format(byte_position, lastbyte)

            part = s3.upload_part_copy(
                Bucket=to_info.bucket,
                Key=to_info.path,
                PartNumber=i,
                UploadId=mpu_id,
                CopySourceRange=srange,
                CopySource={"Bucket": from_info.bucket, "Key": from_info.path},
            )
            parts.append(
                {"PartNumber": i, "ETag": part["CopyPartResult"]["ETag"]}
            )
            byte_position += part_size

        assert n_parts == len(parts)

        s3.complete_multipart_upload(
            Bucket=to_info.bucket,
            Key=to_info.path,
            UploadId=mpu_id,
            MultipartUpload={"Parts": parts},
        )

    @classmethod
    def _copy(cls, s3, from_info, to_info, extra_args):
        # NOTE: object's etag depends on the way it was uploaded to s3 or the
        # way it was copied within the s3. More specifically, it depends on
        # the chunk size that was used to transfer it, which would affect
        # whether an object would be uploaded as a single part or as a
        # multipart.
        #
        # If an object's etag looks like '8978c98bb5a48c2fb5f2c4c905768afa',
        # then it was transferred as a single part, which means that the chunk
        # size used to transfer it was greater or equal to the ContentLength
        # of that object. So to preserve that tag over the next transfer, we
        # could use any value >= ContentLength.
        #
        # If an object's etag looks like '50d67013a5e1a4070bef1fc8eea4d5f9-13',
        # then it was transferred as a multipart, which means that the chunk
        # size used to transfer it was less than ContentLength of that object.
        # Unfortunately, in general, it doesn't mean that the chunk size was
        # the same throughout the transfer, so it means that in order to
        # preserve etag, we need to transfer each part separately, so the
        # object is transfered in the same chunks as it was originally.

        obj = cls.get_head_object(s3, from_info.bucket, from_info.path)
        etag = obj["ETag"].strip('"')
        size = obj["ContentLength"]

        _, _, parts_suffix = etag.partition("-")
        if parts_suffix:
            n_parts = int(parts_suffix)
            cls._copy_multipart(
                s3, from_info, to_info, size, n_parts, extra_args=extra_args
            )
        else:
            source = {"Bucket": from_info.bucket, "Key": from_info.path}
            s3.copy(source, to_info.bucket, to_info.path, ExtraArgs=extra_args)

        cached_etag = cls.get_etag(s3, to_info.bucket, to_info.path)
        if etag != cached_etag:
            raise ETagMismatchError(etag, cached_etag)

    def copy(self, from_info, to_info):
        self._copy(self.s3, from_info, to_info, self.extra_args)

    def remove(self, path_info):
        if path_info.scheme != "s3":
            raise NotImplementedError

        logger.debug("Removing {}".format(path_info))
        self.s3.delete_object(Bucket=path_info.bucket, Key=path_info.path)

    def _list_objects(self, path_info, max_items=None):
        """ Read config for list object api, paginate through list objects."""
        kwargs = {
            "Bucket": path_info.bucket,
            "Prefix": path_info.path,
            "PaginationConfig": {"MaxItems": max_items},
        }
        paginator = self.s3.get_paginator(self.list_objects_api)
        for page in paginator.paginate(**kwargs):
            contents = page.get("Contents", None)
            if not contents:
                continue
            for item in contents:
                yield item

    def _list_paths(self, path_info, max_items=None):
        return (
            item["Key"] for item in self._list_objects(path_info, max_items)
        )

    def list_cache_paths(self):
        return self._list_paths(self.path_info)

    def isfile(self, path_info):
        from botocore.exceptions import ClientError

        if path_info.path.endswith("/"):
            return False

        try:
            self.s3.head_object(Bucket=path_info.bucket, Key=path_info.path)
        except ClientError as exc:
            if exc.response["Error"]["Code"] != "404":
                raise
            return False

        return True

    def exists(self, path_info):
        """Check if the blob exists. If it does not exist,
        it could be a part of a directory path.

        eg: if `data/file.txt` exists, check for `data` should return True
        """
        return self.isfile(path_info) or self.isdir(path_info)

    def makedirs(self, path_info):
        # We need to support creating empty directories, which means
        # creating an object with an empty body and a trailing slash `/`.
        #
        # We are not creating directory objects for every parent prefix,
        # as it is not required.
        dir_path = path_info / ""
        self.s3.put_object(Bucket=path_info.bucket, Key=dir_path.path, Body="")

    def isdir(self, path_info):
        # S3 doesn't have a concept for directories.
        #
        # Using `head_object` with a path pointing to a directory
        # will throw a 404 error.
        #
        # A reliable way to know if a given path is a directory is by
        # checking if there are more files sharing the same prefix
        # with a `list_objects` call.
        #
        # We need to make sure that the path ends with a forward slash,
        # since we can end with false-positives like the following example:
        #
        #       bucket
        #       └── data
        #          ├── alice
        #          └── alpha
        #
        # Using `data/al` as prefix will return `[data/alice, data/alpha]`,
        # While `data/al/` will return nothing.
        #
        dir_path = path_info / ""
        return bool(list(self._list_paths(dir_path, max_items=1)))

    def _upload(self, from_file, to_info, name=None, no_progress_bar=False):
        total = os.path.getsize(from_file)
        with Tqdm(
            disable=no_progress_bar, total=total, bytes=True, desc=name
        ) as pbar:
            self.s3.upload_file(
                from_file,
                to_info.bucket,
                to_info.path,
                Callback=pbar.update,
                ExtraArgs=self.extra_args,
            )

    def _download(self, from_info, to_file, name=None, no_progress_bar=False):
        if no_progress_bar:
            total = None
        else:
            total = self.s3.head_object(
                Bucket=from_info.bucket, Key=from_info.path
            )["ContentLength"]
        with Tqdm(
            disable=no_progress_bar, total=total, bytes=True, desc=name
        ) as pbar:
            self.s3.download_file(
                from_info.bucket, from_info.path, to_file, Callback=pbar.update
            )

    def _generate_download_url(self, path_info, expires=3600):
        params = {"Bucket": path_info.bucket, "Key": path_info.path}
        return self.s3.generate_presigned_url(
            ClientMethod="get_object", Params=params, ExpiresIn=int(expires)
        )

    def walk_files(self, path_info, max_items=None):
        for fname in self._list_paths(path_info / "", max_items):
            if fname.endswith("/"):
                continue

            yield path_info.replace(path=fname)

    def _append_aws_grants_to_extra_args(self, config):
        # Keys for extra_args can be one of the following list:
        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html#boto3.s3.transfer.S3Transfer.ALLOWED_UPLOAD_ARGS
        """
          ALLOWED_UPLOAD_ARGS = [
            'ACL', 'CacheControl', 'ContentDisposition', 'ContentEncoding',
            'ContentLanguage', 'ContentType', 'Expires', 'GrantFullControl',
            'GrantRead', 'GrantReadACP', 'GrantWriteACP', 'Metadata',
            'RequestPayer', 'ServerSideEncryption', 'StorageClass',
            'SSECustomerAlgorithm', 'SSECustomerKey', 'SSECustomerKeyMD5',
            'SSEKMSKeyId', 'WebsiteRedirectLocation'
          ]
        """
        overwriting_acl_grantees = False

        self.grant_full_control = config.get(
            Config.SECTION_AWS_GRANT_FULL_CONTROL, ""
        )
        if self.grant_full_control:
            self.extra_args["GrantFullControl"] = self.grant_full_control
            overwriting_acl_grantees = True

        self.grant_read = config.get(Config.SECTION_AWS_GRANT_READ, "")
        if self.grant_read:
            self.extra_args["GrantRead"] = self.grant_read
            overwriting_acl_grantees = True

        self.grant_read_acp = config.get(Config.SECTION_AWS_GRANT_READ_ACP, "")
        if self.grant_read_acp:
            self.extra_args["GrantReadACP"] = self.grant_read_acp
            overwriting_acl_grantees = True

        self.grant_write_acp = config.get(
            Config.SECTION_AWS_GRANT_WRITE_ACP, ""
        )
        if self.grant_write_acp:
            self.extra_args["GrantWriteACP"] = self.grant_write_acp
            overwriting_acl_grantees = True

        if overwriting_acl_grantees:
            overwriting_fields = [
                Config.SECTION_AWS_GRANT_FULL_CONTROL,
                Config.SECTION_AWS_GRANT_READ,
                Config.SECTION_AWS_GRANT_READ_ACP,
                Config.SECTION_AWS_GRANT_WRITE_ACP,
            ]
            warning_message = [
                "Using one of the [{0}] will".format(
                    ", ".join(overwriting_fields)
                ),
                "overwrite default ACL Grantees list. Default ACL Grantees",
                "list include Owner with FULL_CONTROL permissions.",
            ]
            logger.warning(" ".join(warning_message))
<|endoftext|>


<commit_before>from IPython.core import magic_arguments
from IPython.core.magic import cell_magic, magics_class, Magics
from IPython.utils.path import get_ipython_cache_dir
from IPython.utils import py3compat
import sys
import os
import io
import time
import imp
import json
import subprocess
import shutil
import ast

from mergedict import ConfigDict

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.core import Distribution
from distutils.command.build_ext import build_ext

import pyd.support

@magics_class
class PydMagics(Magics):
    def __init__(self, shell):
        super(PydMagics, self).__init__(shell)
        self._reloads = {}
        self._code_cache = {}

    def _import_all(self, module):
        for k,v in module.__dict__.items():
            if not k.startswith('__'):
                self.shell.push({k:v})

    @magic_arguments.magic_arguments()
    @magic_arguments.argument(
        '-n', '--name',
        help="Specify a name for the Pyd module."
    )
    @magic_arguments.argument(
        '-I', '--include', action='append', default=[],
        help="Add a path to the list of include directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-f', '--force', action='store_true', default=False,
        help="Force the compilation of a new module, even if the source has been "
             "previously compiled."
    )
    @magic_arguments.argument(
        '--compiler', action='store', default='dmd',
        help="Specify the D compiler to be used. Default is dmd"
    )
    @magic_arguments.argument(
        '--compiler_type', action='store', default='dmd',
        help="Specify the compiler type, as in dmd, gdc, ldc or sdc. Needed if "
             "you are using a non-standard compiler name e.g. dmd_HEAD for your"
             "own personal build of dmd from git master HEAD"
    )
    @magic_arguments.argument(
        '--pyd_version', action='store', default='>=0.9.7',
        help="Specify the pyd version to use, as a dub version specifier "
             "see http://code.dlang.org/package-format#version-specs"
    )
    @magic_arguments.argument(
        '--dub_config', action='store', default='{}',
        help='''add your own compilation flags, dependencies etc. as json '''
             '''to be merged with dub.json e.g. {"libs":"fftw3"}. Note '''
             '''that spaces will be interpreted as starting a new argument, '''
             '''If you want spaces in your argument, wrap it in a string'''
    )
    @magic_arguments.argument(
        '--dub_args', action='store', default='',
        help="command line flags to dub. wrap in a string if you want to list "
             "multiple arguments."
    )

    @cell_magic
    def pyd(self, line, cell):

        args = magic_arguments.parse_argstring(self.pyd, line)
        code = 'import ppyd;\n\n\
                extern(C) void PydMain()\n{\n   \
                registerAll!(Alias!(__traits(parent, PydMain)))();\n\
                }\n\n'\
                + cell
        code = code if code.endswith('\n') else code+'\n'

        key = code, line, sys.version_info, sys.executable

        try:
            args.dub_config = json.loads(args.dub_config)
        except:
            args.dub_config = json.loads(ast.literal_eval(args.dub_config))
            pass

        try:
            args.dub_args = ast.literal_eval(args.dub_args)
        except:
            pass

        if args.force:
            # Force a new module name by adding the current time to the
            # key which is hashed to determine the module name.
            key += (time.time(),)
            args.dub_args = '--force ' + args.dub_args

        if args.name:
            module_name = py3compat.unicode_to_str(args.name)
        else:
            module_name = "_pyd_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()

        lib_dir = os.path.join(get_ipython_cache_dir(), 'pyd', module_name)

        if not os.path.exists(lib_dir):
            os.makedirs(lib_dir)

        if os.name == 'nt':
            so_ext = '.dll'
        else:
            so_ext = '.so' #might have to go to dylib on OS X at some point???
        module_path = os.path.join(lib_dir, 'lib' + module_name + so_ext)

        have_module = os.path.isfile(module_path)
        need_pydize = not have_module

        if need_pydize:
            d_include_dirs = args.include
            pyd_file = os.path.join(lib_dir, module_name + '.d')
            pyd_file = py3compat.cast_bytes_py2(pyd_file, encoding=sys.getfilesystemencoding())
            with io.open(pyd_file, 'w', encoding='utf-8') as f:
                f.write(code)

            pyd_dub_file = os.path.join(lib_dir, 'dub.json')
            pyd_dub_file = py3compat.cast_bytes_py2(pyd_dub_file, encoding=sys.getfilesystemencoding())

            pyd_dub_json = json.loads('{}')
            pyd_dub_json['name'] = module_name
            pyd_dub_json['dependencies'] = { "pyd": args.pyd_version, "ppyd": ">=0.1.0" }
            pyd_dub_json['subConfigurations'] = { "pyd": "python{0}{1}".format(sys.version_info.major, sys.version_info.minor) }
            pyd_dub_json['sourceFiles'] = [pyd_file]
            pyd_dub_json['targetType'] = 'dynamicLibrary'
            pyd_dub_json['dflags'] = ['-fPIC']
            pyd_dub_json['libs'] = ['phobos2']

            with io.open(pyd_dub_file, 'w', encoding='utf-8') as f:
                f.write(json.dumps(pyd_dub_json)+'\n')
            try:
                os.remove(os.path.join(lib_dir, 'dub.selections.json'))
            except:
                pass

            dub_desc = json.loads(subprocess.check_output(["dub", "describe", "--root=" + lib_dir], universal_newlines = True))
            for pack in dub_desc['packages']:
                if pack['name'] == 'pyd':
                    _infraDir = os.path.join(pack['path'], 'infrastructure')
                    break

            if os.name == 'nt':
                boilerplatePath = os.path.join(_infraDir, 'd',
                        'python_dll_windows_boilerplate.d'
                        )
            else:
                boilerplatePath = os.path.join(_infraDir, 'd',
                        'python_so_linux_boilerplate.d'
                        )
            pyd_dub_json['sourceFiles'].append(boilerplatePath)

            if args.compiler == 'dmd':
                so_ctor_path = os.path.join(_infraDir, 'd', 'so_ctor.c')
                so_ctor_object_path = os.path.join(lib_dir, "so_ctor.o")
                subprocess.check_call(['cc', "-c", "-fPIC", "-o" + so_ctor_object_path, so_ctor_path])
                pyd_dub_json['sourceFiles'].append(so_ctor_object_path)

            mainTemplate = os.path.join(_infraDir, 'd', 'pydmain_template.d')
            mainTemplateOut = os.path.join(lib_dir, 'pydmain.d')
            with io.open(mainTemplate, 'r', encoding='utf-8') as t, io.open(mainTemplateOut, 'w', encoding='utf-8') as m:
                m.write(t.read() % {'modulename' : module_name})
            pyd_dub_json['sourceFiles'].append(mainTemplateOut)

            pyd_dub_json = ConfigDict(pyd_dub_json)
            pyd_dub_json.merge(args.dub_config)

            with io.open(pyd_dub_file, 'w', encoding='utf-8') as f:
                f.write(json.dumps(pyd_dub_json)+'\n')

            subprocess.check_call(["dub", "-v", "build", "--root=" + lib_dir] + args.dub_args.split(' '))

        if not have_module:
            self._code_cache[key] = module_name

        module = imp.load_dynamic(module_name, module_path)
        self._import_all(module)

def load_ipython_extension(ip):
    ip.register_magics(PydMagics)
<commit_msg>update ppyd dependency<commit_after>from IPython.core import magic_arguments
from IPython.core.magic import cell_magic, magics_class, Magics
from IPython.utils.path import get_ipython_cache_dir
from IPython.utils import py3compat
import sys
import os
import io
import time
import imp
import json
import subprocess
import shutil
import ast

from mergedict import ConfigDict

try:
    import hashlib
except ImportError:
    import md5 as hashlib

from distutils.core import Distribution
from distutils.command.build_ext import build_ext

import pyd.support

@magics_class
class PydMagics(Magics):
    def __init__(self, shell):
        super(PydMagics, self).__init__(shell)
        self._reloads = {}
        self._code_cache = {}

    def _import_all(self, module):
        for k,v in module.__dict__.items():
            if not k.startswith('__'):
                self.shell.push({k:v})

    @magic_arguments.magic_arguments()
    @magic_arguments.argument(
        '-n', '--name',
        help="Specify a name for the Pyd module."
    )
    @magic_arguments.argument(
        '-I', '--include', action='append', default=[],
        help="Add a path to the list of include directories (can be specified "
             "multiple times)."
    )
    @magic_arguments.argument(
        '-f', '--force', action='store_true', default=False,
        help="Force the compilation of a new module, even if the source has been "
             "previously compiled."
    )
    @magic_arguments.argument(
        '--compiler', action='store', default='dmd',
        help="Specify the D compiler to be used. Default is dmd"
    )
    @magic_arguments.argument(
        '--compiler_type', action='store', default='dmd',
        help="Specify the compiler type, as in dmd, gdc, ldc or sdc. Needed if "
             "you are using a non-standard compiler name e.g. dmd_HEAD for your"
             "own personal build of dmd from git master HEAD"
    )
    @magic_arguments.argument(
        '--pyd_version', action='store', default='>=0.9.7',
        help="Specify the pyd version to use, as a dub version specifier "
             "see http://code.dlang.org/package-format#version-specs"
    )
    @magic_arguments.argument(
        '--dub_config', action='store', default='{}',
        help='''add your own compilation flags, dependencies etc. as json '''
             '''to be merged with dub.json e.g. {"libs":"fftw3"}. Note '''
             '''that spaces will be interpreted as starting a new argument, '''
             '''If you want spaces in your argument, wrap it in a string'''
    )
    @magic_arguments.argument(
        '--dub_args', action='store', default='',
        help="command line flags to dub. wrap in a string if you want to list "
             "multiple arguments."
    )

    @cell_magic
    def pyd(self, line, cell):

        args = magic_arguments.parse_argstring(self.pyd, line)
        code = 'import ppyd;\n\n\
                extern(C) void PydMain()\n{\n   \
                registerAll!(Alias!(__traits(parent, PydMain)))();\n\
                }\n\n'\
                + cell
        code = code if code.endswith('\n') else code+'\n'

        key = code, line, sys.version_info, sys.executable

        try:
            args.dub_config = json.loads(args.dub_config)
        except:
            args.dub_config = json.loads(ast.literal_eval(args.dub_config))
            pass

        try:
            args.dub_args = ast.literal_eval(args.dub_args)
        except:
            pass

        if args.force:
            # Force a new module name by adding the current time to the
            # key which is hashed to determine the module name.
            key += (time.time(),)
            args.dub_args = '--force ' + args.dub_args

        if args.name:
            module_name = py3compat.unicode_to_str(args.name)
        else:
            module_name = "_pyd_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest()

        lib_dir = os.path.join(get_ipython_cache_dir(), 'pyd', module_name)

        if not os.path.exists(lib_dir):
            os.makedirs(lib_dir)

        if os.name == 'nt':
            so_ext = '.dll'
        else:
            so_ext = '.so' #might have to go to dylib on OS X at some point???
        module_path = os.path.join(lib_dir, 'lib' + module_name + so_ext)

        have_module = os.path.isfile(module_path)
        need_pydize = not have_module

        if need_pydize:
            d_include_dirs = args.include
            pyd_file = os.path.join(lib_dir, module_name + '.d')
            pyd_file = py3compat.cast_bytes_py2(pyd_file, encoding=sys.getfilesystemencoding())
            with io.open(pyd_file, 'w', encoding='utf-8') as f:
                f.write(code)

            pyd_dub_file = os.path.join(lib_dir, 'dub.json')
            pyd_dub_file = py3compat.cast_bytes_py2(pyd_dub_file, encoding=sys.getfilesystemencoding())

            pyd_dub_json = json.loads('{}')
            pyd_dub_json['name'] = module_name
            pyd_dub_json['dependencies'] = { "pyd": args.pyd_version, "ppyd": ">=0.1.1" }
            pyd_dub_json['subConfigurations'] = { "pyd": "python{0}{1}".format(sys.version_info.major, sys.version_info.minor) }
            pyd_dub_json['sourceFiles'] = [pyd_file]
            pyd_dub_json['targetType'] = 'dynamicLibrary'
            pyd_dub_json['dflags'] = ['-fPIC']
            pyd_dub_json['libs'] = ['phobos2']

            with io.open(pyd_dub_file, 'w', encoding='utf-8') as f:
                f.write(json.dumps(pyd_dub_json)+'\n')
            try:
                os.remove(os.path.join(lib_dir, 'dub.selections.json'))
            except:
                pass

            dub_desc = json.loads(subprocess.check_output(["dub", "describe", "--root=" + lib_dir], universal_newlines = True))
            for pack in dub_desc['packages']:
                if pack['name'] == 'pyd':
                    _infraDir = os.path.join(pack['path'], 'infrastructure')
                    break

            if os.name == 'nt':
                boilerplatePath = os.path.join(_infraDir, 'd',
                        'python_dll_windows_boilerplate.d'
                        )
            else:
                boilerplatePath = os.path.join(_infraDir, 'd',
                        'python_so_linux_boilerplate.d'
                        )
            pyd_dub_json['sourceFiles'].append(boilerplatePath)

            if args.compiler == 'dmd':
                so_ctor_path = os.path.join(_infraDir, 'd', 'so_ctor.c')
                so_ctor_object_path = os.path.join(lib_dir, "so_ctor.o")
                subprocess.check_call(['cc', "-c", "-fPIC", "-o" + so_ctor_object_path, so_ctor_path])
                pyd_dub_json['sourceFiles'].append(so_ctor_object_path)

            mainTemplate = os.path.join(_infraDir, 'd', 'pydmain_template.d')
            mainTemplateOut = os.path.join(lib_dir, 'pydmain.d')
            with io.open(mainTemplate, 'r', encoding='utf-8') as t, io.open(mainTemplateOut, 'w', encoding='utf-8') as m:
                m.write(t.read() % {'modulename' : module_name})
            pyd_dub_json['sourceFiles'].append(mainTemplateOut)

            pyd_dub_json = ConfigDict(pyd_dub_json)
            pyd_dub_json.merge(args.dub_config)

            with io.open(pyd_dub_file, 'w', encoding='utf-8') as f:
                f.write(json.dumps(pyd_dub_json)+'\n')

            subprocess.check_call(["dub", "-v", "build", "--root=" + lib_dir] + args.dub_args.split(' '))

        if not have_module:
            self._code_cache[key] = module_name

        module = imp.load_dynamic(module_name, module_path)
        self._import_all(module)

def load_ipython_extension(ip):
    ip.register_magics(PydMagics)
<|endoftext|>


<commit_before>"""
DSA Signature Algorithm - A simple implementation in Python

Copyright 2016 by Reiner Rottmann <reiner@rottmann.it>

Released under the MIT license.
"""
import os
import random
import hashlib
import subprocess

import time

# Example DSA keypair - real world
dsa_key = {
    'Q': 1218442816993522937915646204915776994404649089503L,
    'P': 11220611807188583130302963536190351192186270126479330588604287699892081267588448305835704397593153801135202051719876685351614175538253684346816652027037363L,
    'G': 11189361631195852088154673407566885728548496486362662112597687161142104619469702160215294558351391466982303919803857229515093575816938371433954759500448775L,
    'pub': 4572510396595314270786423212039255215498677297795049756997099191729339616558419010431226927123876238239229467750410441342637393785565872285607741290303779L,
    'priv': 148102768779017960166999813987055538077373228390L}

# Example DSA keypair (simple) - for humans
dsa_key = {
    'Q': 11,
    'P': 23,
    'G': 4,
    'pub': 8,
    'priv': 7}

def _random_s(min, max):
    """
    Helper function to select a random number.
    :param min: smallest random number
    :param max: largest random number
    :return: random number
    """
    return random.randint(min, max)

def _mod_inverse(a, b):
    """
    Helper function that calculates the Modular multiplicative inverse

    See https://en.wikipedia.org/wiki/Modular_multiplicative_inverse

    Implementation using Extended Euclidean algorithm by Eric taken from:
    http://stackoverflow.com/questions/4798654/modular-multiplicative-inverse-function-in-python

    :param a: integer to calculate the mod inverse from
    :param b: the modulo to use
    :return: Modular multiplicative inverse
    """
    r = -1
    B = b
    A = a
    eq_set = []
    full_set = []
    mod_set = []

    # euclid's algorithm
    while r != 1 and r != 0:
        r = b % a
        q = b // a
        eq_set = [r, b, a, q * -1]
        b = a
        a = r
        full_set.append(eq_set)

    for i in range(0, 4):
        mod_set.append(full_set[-1][i])

    mod_set.insert(2, 1)
    counter = 0

    #extended euclid's algorithm
    for i in range(1, len(full_set)):
        if counter % 2 == 0:
            mod_set[2] = full_set[-1 * (i + 1)][3] * mod_set[4] + mod_set[2]
            mod_set[3] = full_set[-1 * (i + 1)][1]

        elif counter % 2 != 0:
            mod_set[4] = full_set[-1 * (i + 1)][3] * mod_set[2] + mod_set[4]
            mod_set[1] = full_set[-1 * (i + 1)][1]

        counter += 1

    if mod_set[3] == B:
        return mod_set[2] % B
    return mod_set[4] % B

def modexp_lr_k_ary(a, b, n, k=5):
    """
    Compute a ** b (mod n)
        K-ary LR method, with a customizable 'k'.

    This efficient modular exponentiation algorithm was
    implemented by Eli Bendersky

    http://eli.thegreenplace.net/2009/03/28/efficient-modular-exponentiation-algorithms/

    :param a: base
    :param b: exponent
    :param n: modulo
    :param k: customizeable k
    :return: Modular exponentation
    """
    base = 2 << (k - 1)
    # Precompute the table of exponents
    table = [1] * base
    for i in xrange(1, base):
        table[i] = table[i - 1] * a % n
    # Just like the binary LR method, just with a
    # different base
    #
    r = 1
    for digit in reversed(_digits_of_n(b, base)):
        for i in xrange(k):
            r = r * r % n
        if digit:
            r = r * table[digit] % n
    return r

def _digits_of_n(n, b):
    """

    Return the list of the digits in the base 'b'
    representation of n, from LSB to MSB

    This helper function is used by modexp_lr_k_ary and was
    implemented by Eli Bendersky.

    http://eli.thegreenplace.net/2009/03/28/efficient-modular-exponentiation-algorithms/

    :param n: integer
    :param b: base
    :return: number of digits in the base b
    """
    digits = []
    while n:
        digits.append(int(n % b))
        n /= b
    return digits

def dsa_sign(q, p, g, x, message):
    """
    Create a DSA signature of a message
    using the private part of a DSA keypair.

    The message is integer and usually a SHA-1 hash.

    public key: q,p,g, y
    public key: q,p,g, x

    Implemented using code snippets and explanations from:
    * http://www.herongyang.com/Cryptography/DSA-Introduction-Algorithm-Illustration-p23-q11.html
    * https://en.wikipedia.org/wiki/Digital_Signature_Algorithm
    * http://www.docjar.org/html/api/org/bouncycastle/crypto/

    >>> import hashlib
    >>> import dsa
    >>> m = hashlib.sha1()
    >>> m.update("ABCDE")
    >>> message = int("0x" + m.hexdigest(), 0)
    >>> dsa_key = {
    ...     'Q': 11,
    ...     'P': 23,
    ...     'G': 4,
    ...     'pub': 8,
    ...     'priv': 7}
    >>> sig = dsa.dsa_sign(dsa_key["Q"], dsa_key["P"], dsa_key["G"], dsa_key["priv"], message)
    >>> print len(sig)
    2

    :param q: selected prime divisor
    :param p: computed prime modulus: (p-1) mod q = 0
    :param g: computed:
              1 < g < p, g**q mod p = 1
              and
              g = h**((p-1)/q) mod p
    :param x: selected: 0 < x < q
    :param message: message to sign
    :return: DSA signature (s1,s2) sometimes called (r,s)
    """
    s = _random_s(1, q)
    s1 = 0
    s2 = 0
    while True:
        modexp = modexp_lr_k_ary(g, s, p)
        s1 = modexp % q
        if s1 == 0:
            s = _random_s(1, q)
            continue
        s = _mod_inverse(s, q) * (message + x * s1)
        s2 = s % q
        if s2 == 0:
            s = _random_s(1, q)
            continue
        return (int(s1), int(s2))

def dsa_verify(s1, s2, g, p, q, y, message):
    """
    Verify the DSA signature of a message
    using the public part of a DSA keypair.

    The message is integer and usually a SHA-1 hash.

    s1,s2: DSA signature; sometimes called (r,s)

    public key: q,p,g, y
    public key: q,p,g, x

    Implemented using code snippets and explanations from:
    * http://www.herongyang.com/Cryptography/DSA-Introduction-Algorithm-Illustration-p23-q11.html
    * https://en.wikipedia.org/wiki/Digital_Signature_Algorithm
    * http://www.docjar.org/html/api/org/bouncycastle/crypto/

    >>> import hashlib
    >>> import dsa
    >>> m = hashlib.sha1()
    >>> m.update("ABCDE")
    >>> message = int("0x" + m.hexdigest(), 0)
    >>> dsa_key = {
    ...     'Q': 11,
    ...     'P': 23,
    ...     'G': 4,
    ...     'pub': 8,
    ...     'priv': 7}
    >>> sig = (2,3)
    >>> print dsa.dsa_verify(sig[0], sig[1], dsa_key["G"], dsa_key["P"], dsa_key["Q"], dsa_key["pub"], message)
    True

    :param s1: DSA signature part 1, sometimes called r
    :param s2: DSA signature part 2, sometimes called s
    :param q: selected prime divisor
    :param p: computed prime modulus: (p-1) mod q = 0
    :param g: computed:
              1 < g < p, g**q mod p = 1
              and
              g = h**((p-1)/q) mod p
    :param y: computed: y = g**x mod p
    :param message: message to verify
    :return: True or False
    """
    if not s1 > 0:
        return False
    if not s1 < q:
        return False
    if not s2 > 0:
        return False
    if not s2 < q:
        return False
    w = _mod_inverse(s2, q)
    u1 = (message * w) % q
    u2 = (s1 * w) % q
    # v = (((g**u1)*(y**u2)) % p ) % q # correct formula but slooooow!
    # because of that, we use modulo arithmetic to calculate intermediate values:
    u1 = pow(g, u1, p)
    u2 = pow(y, u2, p)
    v = u1 * u2 % p % q
    if v == s1:
        return True
    return False

if __name__ == "__main__":
    import doctest
    doctest.testmod()
    # message = "test"
    # m = hashlib.sha1()
    # m.update(message)
    # sha1 = int("0x" + m.hexdigest(), 0)
    # print sha1
    # # sha1=3 ##delme
    # if dsa_key == None:
    #     dsa_key = dsa_keygen()
    # sig = dsa_sign(dsa_key["Q"], dsa_key["P"], dsa_key["G"], dsa_key["priv"], sha1)
    # #sig = dsa_sign(11, 23, 4, 7, sha1)
    # #print dsa_verify(sig[0], sig[1], 4, 23, 11, 8, 3) #test
    # #s1, s2, g, p, q, pub,  message
    # print dsa_verify(sig[0], sig[1], dsa_key["G"], dsa_key["P"], dsa_key["Q"], dsa_key["pub"], sha1)
<commit_msg>Removing unneeded import.<commit_after>"""
DSA Signature Algorithm - A simple implementation in Python

Copyright 2016 by Reiner Rottmann <reiner@rottmann.it>

Released under the MIT license.
"""
import os
import random
import hashlib
import subprocess

# Example DSA keypair - real world
dsa_key = {
    'Q': 1218442816993522937915646204915776994404649089503L,
    'P': 11220611807188583130302963536190351192186270126479330588604287699892081267588448305835704397593153801135202051719876685351614175538253684346816652027037363L,
    'G': 11189361631195852088154673407566885728548496486362662112597687161142104619469702160215294558351391466982303919803857229515093575816938371433954759500448775L,
    'pub': 4572510396595314270786423212039255215498677297795049756997099191729339616558419010431226927123876238239229467750410441342637393785565872285607741290303779L,
    'priv': 148102768779017960166999813987055538077373228390L}

# Example DSA keypair (simple) - for humans
dsa_key = {
    'Q': 11,
    'P': 23,
    'G': 4,
    'pub': 8,
    'priv': 7}

def _random_s(min, max):
    """
    Helper function to select a random number.
    :param min: smallest random number
    :param max: largest random number
    :return: random number
    """
    return random.randint(min, max)

def _mod_inverse(a, b):
    """
    Helper function that calculates the Modular multiplicative inverse

    See https://en.wikipedia.org/wiki/Modular_multiplicative_inverse

    Implementation using Extended Euclidean algorithm by Eric taken from:
    http://stackoverflow.com/questions/4798654/modular-multiplicative-inverse-function-in-python

    :param a: integer to calculate the mod inverse from
    :param b: the modulo to use
    :return: Modular multiplicative inverse
    """
    r = -1
    B = b
    A = a
    eq_set = []
    full_set = []
    mod_set = []

    # euclid's algorithm
    while r != 1 and r != 0:
        r = b % a
        q = b // a
        eq_set = [r, b, a, q * -1]
        b = a
        a = r
        full_set.append(eq_set)

    for i in range(0, 4):
        mod_set.append(full_set[-1][i])

    mod_set.insert(2, 1)
    counter = 0

    #extended euclid's algorithm
    for i in range(1, len(full_set)):
        if counter % 2 == 0:
            mod_set[2] = full_set[-1 * (i + 1)][3] * mod_set[4] + mod_set[2]
            mod_set[3] = full_set[-1 * (i + 1)][1]

        elif counter % 2 != 0:
            mod_set[4] = full_set[-1 * (i + 1)][3] * mod_set[2] + mod_set[4]
            mod_set[1] = full_set[-1 * (i + 1)][1]

        counter += 1

    if mod_set[3] == B:
        return mod_set[2] % B
    return mod_set[4] % B

def modexp_lr_k_ary(a, b, n, k=5):
    """
    Compute a ** b (mod n)
        K-ary LR method, with a customizable 'k'.

    This efficient modular exponentiation algorithm was
    implemented by Eli Bendersky

    http://eli.thegreenplace.net/2009/03/28/efficient-modular-exponentiation-algorithms/

    :param a: base
    :param b: exponent
    :param n: modulo
    :param k: customizeable k
    :return: Modular exponentation
    """
    base = 2 << (k - 1)
    # Precompute the table of exponents
    table = [1] * base
    for i in xrange(1, base):
        table[i] = table[i - 1] * a % n
    # Just like the binary LR method, just with a
    # different base
    #
    r = 1
    for digit in reversed(_digits_of_n(b, base)):
        for i in xrange(k):
            r = r * r % n
        if digit:
            r = r * table[digit] % n
    return r

def _digits_of_n(n, b):
    """

    Return the list of the digits in the base 'b'
    representation of n, from LSB to MSB

    This helper function is used by modexp_lr_k_ary and was
    implemented by Eli Bendersky.

    http://eli.thegreenplace.net/2009/03/28/efficient-modular-exponentiation-algorithms/

    :param n: integer
    :param b: base
    :return: number of digits in the base b
    """
    digits = []
    while n:
        digits.append(int(n % b))
        n /= b
    return digits

def dsa_sign(q, p, g, x, message):
    """
    Create a DSA signature of a message
    using the private part of a DSA keypair.

    The message is integer and usually a SHA-1 hash.

    public key: q,p,g, y
    public key: q,p,g, x

    Implemented using code snippets and explanations from:
    * http://www.herongyang.com/Cryptography/DSA-Introduction-Algorithm-Illustration-p23-q11.html
    * https://en.wikipedia.org/wiki/Digital_Signature_Algorithm
    * http://www.docjar.org/html/api/org/bouncycastle/crypto/

    >>> import hashlib
    >>> import dsa
    >>> m = hashlib.sha1()
    >>> m.update("ABCDE")
    >>> message = int("0x" + m.hexdigest(), 0)
    >>> dsa_key = {
    ...     'Q': 11,
    ...     'P': 23,
    ...     'G': 4,
    ...     'pub': 8,
    ...     'priv': 7}
    >>> sig = dsa.dsa_sign(dsa_key["Q"], dsa_key["P"], dsa_key["G"], dsa_key["priv"], message)
    >>> print len(sig)
    2

    :param q: selected prime divisor
    :param p: computed prime modulus: (p-1) mod q = 0
    :param g: computed:
              1 < g < p, g**q mod p = 1
              and
              g = h**((p-1)/q) mod p
    :param x: selected: 0 < x < q
    :param message: message to sign
    :return: DSA signature (s1,s2) sometimes called (r,s)
    """
    s = _random_s(1, q)
    s1 = 0
    s2 = 0
    while True:
        modexp = modexp_lr_k_ary(g, s, p)
        s1 = modexp % q
        if s1 == 0:
            s = _random_s(1, q)
            continue
        s = _mod_inverse(s, q) * (message + x * s1)
        s2 = s % q
        if s2 == 0:
            s = _random_s(1, q)
            continue
        return (int(s1), int(s2))

def dsa_verify(s1, s2, g, p, q, y, message):
    """
    Verify the DSA signature of a message
    using the public part of a DSA keypair.

    The message is integer and usually a SHA-1 hash.

    s1,s2: DSA signature; sometimes called (r,s)

    public key: q,p,g, y
    public key: q,p,g, x

    Implemented using code snippets and explanations from:
    * http://www.herongyang.com/Cryptography/DSA-Introduction-Algorithm-Illustration-p23-q11.html
    * https://en.wikipedia.org/wiki/Digital_Signature_Algorithm
    * http://www.docjar.org/html/api/org/bouncycastle/crypto/

    >>> import hashlib
    >>> import dsa
    >>> m = hashlib.sha1()
    >>> m.update("ABCDE")
    >>> message = int("0x" + m.hexdigest(), 0)
    >>> dsa_key = {
    ...     'Q': 11,
    ...     'P': 23,
    ...     'G': 4,
    ...     'pub': 8,
    ...     'priv': 7}
    >>> sig = (2,3)
    >>> print dsa.dsa_verify(sig[0], sig[1], dsa_key["G"], dsa_key["P"], dsa_key["Q"], dsa_key["pub"], message)
    True

    :param s1: DSA signature part 1, sometimes called r
    :param s2: DSA signature part 2, sometimes called s
    :param q: selected prime divisor
    :param p: computed prime modulus: (p-1) mod q = 0
    :param g: computed:
              1 < g < p, g**q mod p = 1
              and
              g = h**((p-1)/q) mod p
    :param y: computed: y = g**x mod p
    :param message: message to verify
    :return: True or False
    """
    if not s1 > 0:
        return False
    if not s1 < q:
        return False
    if not s2 > 0:
        return False
    if not s2 < q:
        return False
    w = _mod_inverse(s2, q)
    u1 = (message * w) % q
    u2 = (s1 * w) % q
    # v = (((g**u1)*(y**u2)) % p ) % q # correct formula but slooooow!
    # because of that, we use modulo arithmetic to calculate intermediate values:
    u1 = pow(g, u1, p)
    u2 = pow(y, u2, p)
    v = u1 * u2 % p % q
    if v == s1:
        return True
    return False

if __name__ == "__main__":
    import doctest
    doctest.testmod()
    # message = "test"
    # m = hashlib.sha1()
    # m.update(message)
    # sha1 = int("0x" + m.hexdigest(), 0)
    # print sha1
    # # sha1=3 ##delme
    # if dsa_key == None:
    #     dsa_key = dsa_keygen()
    # sig = dsa_sign(dsa_key["Q"], dsa_key["P"], dsa_key["G"], dsa_key["priv"], sha1)
    # #sig = dsa_sign(11, 23, 4, 7, sha1)
    # #print dsa_verify(sig[0], sig[1], 4, 23, 11, 8, 3) #test
    # #s1, s2, g, p, q, pub,  message
    # print dsa_verify(sig[0], sig[1], dsa_key["G"], dsa_key["P"], dsa_key["Q"], dsa_key["pub"], sha1)
<|endoftext|>


<commit_before>from .utils import EqualityComparableUsingAttributeDictionary
from .filters import Filter, TermFilter, TermsFilter, ANDFilter, NotFilter

class AggFactory(EqualityComparableUsingAttributeDictionary):

    def __init__(self):
        self.aggs = []

    def add(self, agg):
        """Add a term factory"""
        self.aggs.append(agg)

    def reset(self):
        """Reset the aggs"""
        self.aggs = []

    def serialize(self):
        res = {}
        for agg in self.aggs:
            res.update(agg.serialize())
        return res

class Agg(EqualityComparableUsingAttributeDictionary):

    def __init__(self, name, scope=None, nested=None, is_global=None,
                 agg_filter=None, **kwargs):
        self.name = name
        self.scope = scope
        self.nested = nested
        self.is_global = is_global
        self.agg_filter = agg_filter

    def serialize(self):
        data = {self._internal_name: self._serialize()}
        if self.scope is not None:
            data["scope"] = self.scope
        if self.nested is not None:
            data["nested"] = self.nested
        if self.is_global:
            data['global'] = self.is_global
        if self.agg_filter:
            data['agg_filter'] = self.agg_filter.serialize()
        if isinstance(self, BucketAgg):
            return {self.name: data}
        return {self.name: data}

    def _serialize(self):
        raise NotImplementedError

    @property
    def _internal_name(self):
        raise NotImplementedError

    @property
    def _name(self):
        return self.name

class BucketAgg(Agg):
    def __init__(self, name, sub_aggs=None, **kwargs):
        super(BucketAgg, self).__init__(name, **kwargs)
        self.sub_aggs = sub_aggs
        self.name = name

    def serialize(self):
        data = super(BucketAgg, self).serialize()
        sub_data = {}
        if self.sub_aggs is not None:
            for sub_agg in self.sub_aggs:
                if isinstance(sub_agg, Agg):
                    sub_data.update(sub_agg.serialize())
                else:
                    raise RuntimeError("Invalid Agg: Only Stats-Aggregations allowed as Sub-Aggregations")
            data[self.name].update({"aggs": sub_data})
        return data

    @property
    def _internal_name(self):
        raise NotImplementedError

class FilterAgg(BucketAgg):

    _internal_name = "filter"

    def __init__(self, name, filter, **kwargs):
        super(FilterAgg, self).__init__(name, **kwargs)
        self.filter = filter

    def _serialize(self):
        return self.filter.serialize()

class HistogramAgg(BucketAgg):

    _internal_name = "histogram"

    def __init__(self, name, field=None, interval=None, time_interval=None,
                 key_field=None, value_field=None, key_script=None,
                 value_script=None, params=None, **kwargs):
        super(HistogramAgg, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_interval = time_interval
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _add_interval(self, data):
        if self.interval:
            data['interval'] = self.interval
        elif self.time_interval:
            data['time_interval'] = self.time_interval
        else:
            raise RuntimeError("Invalid field: interval or time_interval required")

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
            self._add_interval(data)
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            else:
                raise RuntimeError("Invalid key_field: value_field required")
            self._add_interval(data)
        elif self.key_script:
            data['key_script'] = self.key_script
            if self.value_script:
                data['value_script'] = self.value_script
            else:
                raise RuntimeError("Invalid key_script: value_script required")
            if self.params:
                data['params'] = self.params
            if self.interval:
                data['interval'] = self.interval
            elif self.time_interval:
                data['time_interval'] = self.time_interval
        return data

class DateHistogramAgg(BucketAgg):

    _internal_name = "date_histogram"

    def __init__(self, name, field=None, interval=None, time_zone=None, pre_zone=None,
                 post_zone=None, factor=None, pre_offset=None, post_offset=None,
                 key_field=None, value_field=None, value_script=None, params=None, **kwargs):
        super(DateHistogramAgg, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_zone = time_zone
        self.pre_zone = pre_zone
        self.post_zone = post_zone
        self.factor = factor
        self.pre_offset = pre_offset
        self.post_offset = post_offset
        self.key_field = key_field
        self.value_field = value_field
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if self.interval:
            data['interval'] = self.interval
        else:
            raise RuntimeError("interval required")
        if self.time_zone:
            data['time_zone'] = self.time_zone
        if self.pre_zone:
            data['pre_zone'] = self.pre_zone
        if self.post_zone:
            data['post_zone'] = self.post_zone
        if self.factor:
            data['factor'] = self.factor
        if self.pre_offset:
            data['pre_offset'] = self.pre_offset
        if self.post_offset:
            data['post_offset'] = self.post_offset
        if self.field:
            data['field'] = self.field
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            elif self.value_script:
                data['value_script'] = self.value_script
                if self.params:
                    data['params'] = self.params
            else:
                raise RuntimeError("Invalid key_field: value_field or value_script required")
        return data

class RangeAgg(BucketAgg):

    _internal_name = "range"

    def __init__(self, name, field=None, ranges=None, key_field=None, value_field=None,
                 key_script=None, value_script=None, params=None, **kwargs):
        super(RangeAgg, self).__init__(name, **kwargs)
        self.field = field
        self.ranges = ranges or []
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if not self.ranges:
            raise RuntimeError("Invalid ranges")
        data['ranges'] = self.ranges
        if self.field:
            data['field'] = self.field
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            else:
                raise RuntimeError("Invalid key_field: value_field required")
        elif self.key_script:
            data['key_script'] = self.key_script
            if self.value_script:
                data['value_script'] = self.value_script
            else:
                raise RuntimeError("Invalid key_script: value_script required")
            if self.params:
                data['params'] = self.params
        return data

class StatsAgg(Agg):

    _internal_name = "stats"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(StatsAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class ValueCountAgg(Agg):

    _internal_name = "value_count"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(ValueCountAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class SumAgg(Agg):

    _internal_name = "sum"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(SumAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class AvgAgg(Agg):

    _internal_name = "avg"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(AvgAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class TermsAgg(BucketAgg):

    _internal_name = "terms"

    def __init__(self, name, field=None, fields=None, size=100, order=None,
                 exclude=None, regex=None, regex_flags="DOTALL", script=None,
                 lang=None, all_terms=None, **kwargs):
        super(TermsAgg, self).__init__(name, **kwargs)
        self.field = field
        self.fields = fields
        self.size = size
        self.order = order
        self.exclude = exclude or []
        self.regex = regex
        self.regex_flags = regex_flags
        self.script = script
        self.lang = lang
        self.all_terms = all_terms

    def _serialize(self):
        if not self.fields and not self.field and not self.script:
            raise RuntimeError("Field, Fields or Script is required:%s" % self.order)

        data = {}
        if self.fields:
            data['fields'] = self.fields
        elif self.field:
            data['field'] = self.field

        if self.script:
            data['script'] = self.script
            if self.lang:
                data['lang'] = self.lang
        if self.size is not None:
            data['size'] = self.size
        if self.order:
            if self.order not in ['count', 'term', 'reverse_count', 'reverse_term']:
                raise RuntimeError("Invalid order value:%s" % self.order)
            data['order'] = self.order
        if self.exclude:
            data['exclude'] = self.exclude
        if (self.fields or self.field) and self.regex:
            data['regex'] = self.regex
            if self.regex_flags:
                data['regex_flags'] = self.regex_flags
        if self.all_terms:
            data['all_terms'] = self.all_terms
        return data

class TermStatsAgg(Agg):

    _internal_name = "terms_stats"

    def __init__(self, name, size=10, order=None, key_field=None, value_field=None,
                 key_script=None, value_script=None, params=None, **kwargs):
        super(TermStatsAgg, self).__init__(name, **kwargs)
        self.size = size
        self.ORDER_VALUES = ['term', 'reverse_term', 'count', 'reverse_count',
                             'total', 'reverse_total', 'min', 'reverse_min',
                             'max', 'reverse_max', 'mean', 'reverse_mean']
        self.order = order if order is not None else self.ORDER_VALUES[0]
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if self.size is not None:
            data['size'] = self.size
        if self.order:
            if self.order not in self.ORDER_VALUES:
                raise RuntimeError("Invalid order value:%s" % self.order)
            data['order'] = self.order

        if self.key_field:
            data['key_field'] = self.key_field
        else:
            raise RuntimeError("key_field required")

        if self.value_field:
            data['value_field'] = self.value_field
        elif self.value_script:
            data['value_script'] = self.value_script
            if self.params:
                data['params'] = self.params
        else:
            raise RuntimeError("Invalid value: value_field OR value_script required")

        return data

class AggQueryWrap(EqualityComparableUsingAttributeDictionary):

    def __init__(self, wrap_object, **kwargs):
        """Base Object for every Filter Object"""
        self.wrap_object = wrap_object

    def serialize(self):
        return {"query": self.wrap_object.serialize()}

<commit_msg>Added NestedAgg<commit_after>from .utils import EqualityComparableUsingAttributeDictionary
from .filters import Filter, TermFilter, TermsFilter, ANDFilter, NotFilter

class AggFactory(EqualityComparableUsingAttributeDictionary):

    def __init__(self):
        self.aggs = []

    def add(self, agg):
        """Add a term factory"""
        self.aggs.append(agg)

    def reset(self):
        """Reset the aggs"""
        self.aggs = []

    def serialize(self):
        res = {}
        for agg in self.aggs:
            res.update(agg.serialize())
        return res

class Agg(EqualityComparableUsingAttributeDictionary):

    def __init__(self, name, scope=None, nested=None, is_global=None,
                 agg_filter=None, **kwargs):
        self.name = name
        self.scope = scope
        self.nested = nested
        self.is_global = is_global
        self.agg_filter = agg_filter

    def serialize(self):
        data = {self._internal_name: self._serialize()}
        if self.scope is not None:
            data["scope"] = self.scope
        if self.nested is not None:
            data["nested"] = self.nested
        if self.is_global:
            data['global'] = self.is_global
        if self.agg_filter:
            data['agg_filter'] = self.agg_filter.serialize()
        if isinstance(self, BucketAgg):
            return {self.name: data}
        return {self.name: data}

    def _serialize(self):
        raise NotImplementedError

    @property
    def _internal_name(self):
        raise NotImplementedError

    @property
    def _name(self):
        return self.name

class BucketAgg(Agg):
    def __init__(self, name, sub_aggs=None, **kwargs):
        super(BucketAgg, self).__init__(name, **kwargs)
        self.sub_aggs = sub_aggs
        self.name = name

    def serialize(self):
        data = super(BucketAgg, self).serialize()
        sub_data = {}
        if self.sub_aggs is not None:
            for sub_agg in self.sub_aggs:
                if isinstance(sub_agg, Agg):
                    sub_data.update(sub_agg.serialize())
                else:
                    raise RuntimeError("Invalid Agg: Only Stats-Aggregations allowed as Sub-Aggregations")
            data[self.name].update({"aggs": sub_data})
        return data

    @property
    def _internal_name(self):
        raise NotImplementedError

class FilterAgg(BucketAgg):

    _internal_name = "filter"

    def __init__(self, name, filter, **kwargs):
        super(FilterAgg, self).__init__(name, **kwargs)
        self.filter = filter

    def _serialize(self):
        return self.filter.serialize()

class HistogramAgg(BucketAgg):

    _internal_name = "histogram"

    def __init__(self, name, field=None, interval=None, time_interval=None,
                 key_field=None, value_field=None, key_script=None,
                 value_script=None, params=None, **kwargs):
        super(HistogramAgg, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_interval = time_interval
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _add_interval(self, data):
        if self.interval:
            data['interval'] = self.interval
        elif self.time_interval:
            data['time_interval'] = self.time_interval
        else:
            raise RuntimeError("Invalid field: interval or time_interval required")

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
            self._add_interval(data)
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            else:
                raise RuntimeError("Invalid key_field: value_field required")
            self._add_interval(data)
        elif self.key_script:
            data['key_script'] = self.key_script
            if self.value_script:
                data['value_script'] = self.value_script
            else:
                raise RuntimeError("Invalid key_script: value_script required")
            if self.params:
                data['params'] = self.params
            if self.interval:
                data['interval'] = self.interval
            elif self.time_interval:
                data['time_interval'] = self.time_interval
        return data

class DateHistogramAgg(BucketAgg):

    _internal_name = "date_histogram"

    def __init__(self, name, field=None, interval=None, time_zone=None, pre_zone=None,
                 post_zone=None, factor=None, pre_offset=None, post_offset=None,
                 key_field=None, value_field=None, value_script=None, params=None, **kwargs):
        super(DateHistogramAgg, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_zone = time_zone
        self.pre_zone = pre_zone
        self.post_zone = post_zone
        self.factor = factor
        self.pre_offset = pre_offset
        self.post_offset = post_offset
        self.key_field = key_field
        self.value_field = value_field
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if self.interval:
            data['interval'] = self.interval
        else:
            raise RuntimeError("interval required")
        if self.time_zone:
            data['time_zone'] = self.time_zone
        if self.pre_zone:
            data['pre_zone'] = self.pre_zone
        if self.post_zone:
            data['post_zone'] = self.post_zone
        if self.factor:
            data['factor'] = self.factor
        if self.pre_offset:
            data['pre_offset'] = self.pre_offset
        if self.post_offset:
            data['post_offset'] = self.post_offset
        if self.field:
            data['field'] = self.field
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            elif self.value_script:
                data['value_script'] = self.value_script
                if self.params:
                    data['params'] = self.params
            else:
                raise RuntimeError("Invalid key_field: value_field or value_script required")
        return data

class NestedAgg(BucketAgg):
    _internal_name = "nested"

    def __init__(self, name, path, **kwargs):
        super(NestedAgg, self).__init__(name, **kwargs)
        self.path = path

    def _serialize(self):
        data = {}
        data['path'] = self.path
        return data

class RangeAgg(BucketAgg):

    _internal_name = "range"

    def __init__(self, name, field=None, ranges=None, key_field=None, value_field=None,
                 key_script=None, value_script=None, params=None, **kwargs):
        super(RangeAgg, self).__init__(name, **kwargs)
        self.field = field
        self.ranges = ranges or []
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if not self.ranges:
            raise RuntimeError("Invalid ranges")
        data['ranges'] = self.ranges
        if self.field:
            data['field'] = self.field
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            else:
                raise RuntimeError("Invalid key_field: value_field required")
        elif self.key_script:
            data['key_script'] = self.key_script
            if self.value_script:
                data['value_script'] = self.value_script
            else:
                raise RuntimeError("Invalid key_script: value_script required")
            if self.params:
                data['params'] = self.params
        return data

class StatsAgg(Agg):

    _internal_name = "stats"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(StatsAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class ValueCountAgg(Agg):

    _internal_name = "value_count"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(ValueCountAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class SumAgg(Agg):

    _internal_name = "sum"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(SumAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class AvgAgg(Agg):

    _internal_name = "avg"

    def __init__(self, name, field=None, script=None, params=None, **kwargs):
        super(AvgAgg, self).__init__(name, **kwargs)
        self.field = field
        self.script = script
        self.params = params

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
        elif self.script:
            data['script'] = self.script
            if self.params:
                data['params'] = self.params
        return data

class TermsAgg(BucketAgg):

    _internal_name = "terms"

    def __init__(self, name, field=None, fields=None, size=100, order=None,
                 exclude=None, regex=None, regex_flags="DOTALL", script=None,
                 lang=None, all_terms=None, **kwargs):
        super(TermsAgg, self).__init__(name, **kwargs)
        self.field = field
        self.fields = fields
        self.size = size
        self.order = order
        self.exclude = exclude or []
        self.regex = regex
        self.regex_flags = regex_flags
        self.script = script
        self.lang = lang
        self.all_terms = all_terms

    def _serialize(self):
        if not self.fields and not self.field and not self.script:
            raise RuntimeError("Field, Fields or Script is required:%s" % self.order)

        data = {}
        if self.fields:
            data['fields'] = self.fields
        elif self.field:
            data['field'] = self.field

        if self.script:
            data['script'] = self.script
            if self.lang:
                data['lang'] = self.lang
        if self.size is not None:
            data['size'] = self.size
        if self.order:
            if self.order not in ['count', 'term', 'reverse_count', 'reverse_term']:
                raise RuntimeError("Invalid order value:%s" % self.order)
            data['order'] = self.order
        if self.exclude:
            data['exclude'] = self.exclude
        if (self.fields or self.field) and self.regex:
            data['regex'] = self.regex
            if self.regex_flags:
                data['regex_flags'] = self.regex_flags
        if self.all_terms:
            data['all_terms'] = self.all_terms
        return data

class TermStatsAgg(Agg):

    _internal_name = "terms_stats"

    def __init__(self, name, size=10, order=None, key_field=None, value_field=None,
                 key_script=None, value_script=None, params=None, **kwargs):
        super(TermStatsAgg, self).__init__(name, **kwargs)
        self.size = size
        self.ORDER_VALUES = ['term', 'reverse_term', 'count', 'reverse_count',
                             'total', 'reverse_total', 'min', 'reverse_min',
                             'max', 'reverse_max', 'mean', 'reverse_mean']
        self.order = order if order is not None else self.ORDER_VALUES[0]
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if self.size is not None:
            data['size'] = self.size
        if self.order:
            if self.order not in self.ORDER_VALUES:
                raise RuntimeError("Invalid order value:%s" % self.order)
            data['order'] = self.order

        if self.key_field:
            data['key_field'] = self.key_field
        else:
            raise RuntimeError("key_field required")

        if self.value_field:
            data['value_field'] = self.value_field
        elif self.value_script:
            data['value_script'] = self.value_script
            if self.params:
                data['params'] = self.params
        else:
            raise RuntimeError("Invalid value: value_field OR value_script required")

        return data

class AggQueryWrap(EqualityComparableUsingAttributeDictionary):

    def __init__(self, wrap_object, **kwargs):
        """Base Object for every Filter Object"""
        self.wrap_object = wrap_object

    def serialize(self):
        return {"query": self.wrap_object.serialize()}

<|endoftext|>


<commit_before>"""
PyMzn wraps the MiniZinc tools by executing them with the ``subprocess.run``
function from the standard library. The process of running a executable,
checking for errors and returning the output is carried out by the
``pymzn.bin.run`` function, which takes as input either a string or a list of
arguments to pass to the ``pymzn.bin.cmd`` function.

The ``pymzn.bin.cmd`` can be used in this way:
::

    pymzn.bin.cmd('path/to/command', [5, '-f', '--flag2', ('--arg1', 'val1'), ('--arg2', 2)])

which will become:
::

    'path/to/command 5 -f --flag2 --arg1 val1 --arg2 2'

"""
import os
import time
import signal
import logging
import numbers
import subprocess
import collections.abc

def cmd(path, args):
    """
    Returns the command string from the path to the binary with the provided
    arguments.

    :param string path: The path to the binary file to be executed. If the
                        binary is in the PATH, then only the name is needed
    :param list args: A list of arguments to pass to the command. Supported
                      types of arguments are str, int, float and key-value
                      tuples
    :return: The string command
    :rtype: str
    """
    _cmd = [path]
    for arg in args:
        if isinstance(arg, str):
            _cmd.append(arg)
        elif isinstance(arg, numbers.Number):
            _cmd.append(str(arg))
        elif isinstance(arg, collections.abc.Iterable) and len(arg) == 2:
            k, v = arg
            if isinstance(k, str) and isinstance(v, (str, numbers.Number)):
                _cmd.append(str(k))
                _cmd.append(str(v))
            else:
                raise ValueError('Invalid argument: {}'.format(arg))
        else:
            raise TypeError('Invalid argument: {}'.format(arg))
    return ' '.join(_cmd)

def run(arg, stdin=None, timeout=None):
    """
    Executes a shell command and waits for the result.

    :param str arg: The command string to be executed, as returned from the
                    cmd function.
    :param str stdin: String containing the input stream to pass to
                      the command
    :return: A string containing the output stream of the command
    :rtype: str
    """
    log = logging.getLogger(__name__)

    if isinstance(arg, list):
        arg = cmd(arg[0], arg[1:])

    log.debug('Executing command: %s', arg, extra={'stdin': stdin})
    start = time.time()
    proc = subprocess.Popen(arg, shell=True, bufsize=1,
                          universal_newlines=True,
                          stdin=subprocess.PIPE,
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE,
                          preexec_fn=os.setsid)
    try:
        out, err = proc.communicate(stdin, timeout=timeout)
        ret = proc.wait()
        if ret:
            raise RuntimeError(err)
    except subprocess.TimeoutExpired:
        proc.kill()
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        out, err = proc.communicate()
    end = time.time()
    log.debug('Done. Running time: {0:.2f} seconds'.format(end - start))
    return out

def stream(arg, stdin=None):
    """
    Executes a shell command and generates lines of output without waiting
    for it to finish.

    :param str arg: The command string to be executed, as returned from the
                    cmd function.
    :param str stdin: Input stream to pass to the command
    :return: A generator containing the lines in the output stream
             of the command
    :rtype: generator of str
    """
    log = logging.getLogger(__name__)
    log.debug('Executing streaming command: %s', arg, extra={'stdin': stdin})
    proc = subprocess.Popen(arg, shell=True, bufsize=1,
                            universal_newlines=True,
                            stdin=subprocess.PIPE,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    if stdin:
        proc.stdin.write(stdin)
        proc.stdin.close()

    out = []
    for line in proc.stdout:
        out.append(line)
        yield line
    out = ''.join(out)
    err = proc.stderr.read()

    ret = proc.wait()
    if ret:
        raise subprocess.CalledProcessError(ret, arg, out, err)
<commit_msg>bin: adapt to new logging<commit_after>"""
PyMzn wraps the MiniZinc tools by executing them with the ``subprocess.run``
function from the standard library. The process of running a executable,
checking for errors and returning the output is carried out by the
``pymzn.bin.run`` function, which takes as input either a string or a list of
arguments to pass to the ``pymzn.bin.cmd`` function.

The ``pymzn.bin.cmd`` can be used in this way:
::

    pymzn.bin.cmd('path/to/command', [5, '-f', '--flag2', ('--arg1', 'val1'), ('--arg2', 2)])

which will become:
::

    'path/to/command 5 -f --flag2 --arg1 val1 --arg2 2'

"""
import os
import time
import signal
import numbers
import subprocess
import collections.abc

from ._utils import get_logger

def cmd(path, args):
    """
    Returns the command string from the path to the binary with the provided
    arguments.

    :param string path: The path to the binary file to be executed. If the
                        binary is in the PATH, then only the name is needed
    :param list args: A list of arguments to pass to the command. Supported
                      types of arguments are str, int, float and key-value
                      tuples
    :return: The string command
    :rtype: str
    """
    _cmd = [path]
    for arg in args:
        if isinstance(arg, str):
            _cmd.append(arg)
        elif isinstance(arg, numbers.Number):
            _cmd.append(str(arg))
        elif isinstance(arg, collections.abc.Iterable) and len(arg) == 2:
            k, v = arg
            if isinstance(k, str) and isinstance(v, (str, numbers.Number)):
                _cmd.append(str(k))
                _cmd.append(str(v))
            else:
                raise ValueError('Invalid argument: {}'.format(arg))
        else:
            raise TypeError('Invalid argument: {}'.format(arg))
    return ' '.join(_cmd)

def run(arg, stdin=None, timeout=None):
    """
    Executes a shell command and waits for the result.

    :param str arg: The command string to be executed, as returned from the
                    cmd function.
    :param str stdin: String containing the input stream to pass to
                      the command
    :return: A string containing the output stream of the command
    :rtype: str
    """
    log = get_logger(__name__)

    if isinstance(arg, list):
        arg = cmd(arg[0], arg[1:])

    log.debug('Executing command: %s', arg, extra={'stdin': stdin})
    start = time.time()
    proc = subprocess.Popen(arg, shell=True, bufsize=1,
                            universal_newlines=True,
                            stdin=subprocess.PIPE,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            preexec_fn=os.setsid)
    try:
        out, err = proc.communicate(stdin, timeout=timeout)
        ret = proc.wait()
        if ret:
            raise RuntimeError(err)
    except subprocess.TimeoutExpired:
        proc.kill()
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        out, err = proc.communicate()
    end = time.time()
    log.debug('Done. Running time: {0:.2f} seconds'.format(end - start))
    return out

def stream(arg, stdin=None):
    """
    Executes a shell command and generates lines of output without waiting
    for it to finish.

    :param str arg: The command string to be executed, as returned from the
                    cmd function.
    :param str stdin: Input stream to pass to the command
    :return: A generator containing the lines in the output stream
             of the command
    :rtype: generator of str
    """
    log = get_logger(__name__)
    log.debug('Executing streaming command: %s', arg, extra={'stdin': stdin})
    proc = subprocess.Popen(arg, shell=True, bufsize=1,
                            universal_newlines=True,
                            stdin=subprocess.PIPE,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    if stdin:
        proc.stdin.write(stdin)
        proc.stdin.close()

    out = []
    for line in proc.stdout:
        out.append(line)
        yield line
    out = ''.join(out)
    err = proc.stderr.read()

    ret = proc.wait()
    if ret:
        raise subprocess.CalledProcessError(ret, arg, out, err)
<|endoftext|>


<commit_before># Python Security Project (PySec) and its related class files.
#
# PySec is a set of tools for secure application development under Linux
#
# Copyright 2014 PySec development team
#
# 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.
#
# -*- coding: ascii -*-
from itertools import islice

def knp(source, pattern, start, stop):
    """Yields all oocurrencies of pattern in source[start:stop]"""
    shifts = [1] * (len(pattern) + 1)
    shift = 1
    for pos in range(len(pattern)):
        while pattern[pos] != pattern[pos - shift] and shift <= pos:
            shift += shifts[pos - shift]
        shifts[pos + 1] = shift
    # search pattern
    mlen = 0
    plen = len(pattern)
    for sub in islice(source, start, stop):
        while mlen == plen or mlen >= 0 and pattern[mlen] != sub:
            sl = shifts[mlen]
            start += sl
            mlen -= sl
        mlen += 1
        if mlen == plen:
            yield start

def knp_find(source, pattern, start, stop):
    """Returns a true value if find the pattern in source"""
    for index in knp(source, pattern, start, stop):
        return 1
    return 0
<commit_msg>add knp_first function<commit_after># Python Security Project (PySec) and its related class files.
#
# PySec is a set of tools for secure application development under Linux
#
# Copyright 2014 PySec development team
#
# 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.
#
# -*- coding: ascii -*-
from itertools import islice

def knp(source, pattern, start, stop):
    """Yields all oocurrencies of pattern in source[start:stop]"""
    shifts = [1] * (len(pattern) + 1)
    shift = 1
    for pos in range(len(pattern)):
        while pattern[pos] != pattern[pos - shift] and shift <= pos:
            shift += shifts[pos - shift]
        shifts[pos + 1] = shift
    # search pattern
    mlen = 0
    plen = len(pattern)
    for sub in islice(source, start, stop):
        while mlen == plen or mlen >= 0 and pattern[mlen] != sub:
            sl = shifts[mlen]
            start += sl
            mlen -= sl
        mlen += 1
        if mlen == plen:
            yield start

def knp_first(source, pattern, start, stop):
    # XXX check if duplicate knp() code is much faster than generator creation
    try:
        return knp(source, pattern, start, stop).next()
    except StopIteration:
        return -1

def knp_find(source, pattern, start, stop):
    """Returns a true value if find the pattern in source"""
    for index in knp(source, pattern, start, stop):
        return 1
    return 0
<|endoftext|>


<commit_before><commit_msg>[bugfix] __uses_HD17 should always be defined<commit_after><|endoftext|>