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

# (c) 2014, Sebastien Rohaut <sebastien.rohaut@gmail.com>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function
__metaclass__ = type

ANSIBLE_METADATA = {'metadata_version': '1.1',
                    'status': ['preview'],
                    'supported_by': 'community'}

DOCUMENTATION = '''
---
module: pam_limits
version_added: "2.0"
author:
    - "Sebastien Rohaut (@usawa)"
short_description: Modify Linux PAM limits
description:
     - The C(pam_limits) module modify PAM limits, default in /etc/security/limits.conf.
       For the full documentation, see man limits.conf(5).
options:
  domain:
    description:
      - A username, @groupname, wildcard, uid/gid range.
    required: true
  limit_type:
    description:
      - Limit type, see C(man limits) for an explanation
    required: true
    choices: [ "hard", "soft", "-" ]
  limit_item:
    description:
      - The limit to be set
    required: true
    choices:
        - "core"
        - "data"
        - "fsize"
        - "memlock"
        - "nofile"
        - "rss"
        - "stack"
        - "cpu"
        - "nproc"
        - "as"
        - "maxlogins"
        - "maxsyslogins"
        - "priority"
        - "locks"
        - "sigpending"
        - "msgqueue"
        - "nice"
        - "rtprio"
        - "chroot"
  value:
    description:
      - The value of the limit.
    required: true
  backup:
    description:
      - Create a backup file including the timestamp information so you can get
        the original file back if you somehow clobbered it incorrectly.
    required: false
    choices: [ "yes", "no" ]
    default: "no"
  use_min:
    description:
      - If set to C(yes), the minimal value will be used or conserved.
        If the specified value is inferior to the value in the file, file content is replaced with the new value,
        else content is not modified.
    required: false
    choices: [ "yes", "no" ]
    default: "no"
  use_max:
    description:
      - If set to C(yes), the maximal value will be used or conserved.
        If the specified value is superior to the value in the file, file content is replaced with the new value,
        else content is not modified.
    required: false
    choices: [ "yes", "no" ]
    default: "no"
  dest:
    description:
      - Modify the limits.conf path.
    required: false
    default: "/etc/security/limits.conf"
  comment:
    description:
      - Comment associated with the limit.
    required: false
    default: ''
'''

EXAMPLES = '''
# Add or modify nofile soft limit for the user joe
- pam_limits:
    domain: joe
    limit_type: soft
    limit_item: nofile
    value: 64000

# Add or modify fsize hard limit for the user smith. Keep or set the maximal value.
- pam_limits:
    domain: smith
    limit_type: hard
    limit_item: fsize
    value: 1000000
    use_max: yes

# Add or modify memlock, both soft and hard, limit for the user james with a comment.
- pam_limits:
    domain: james
    limit_type: '-'
    limit_item: memlock
    value: unlimited
    comment: unlimited memory lock for james
'''

import os
import os.path
import tempfile
import re

from ansible.module_utils.basic import AnsibleModule
from ansible.module_utils._text import to_native

def main():

    pam_items = ['core', 'data', 'fsize', 'memlock', 'nofile', 'rss', 'stack', 'cpu', 'nproc', 'as', 'maxlogins', 'maxsyslogins', 'priority', 'locks',
                 'sigpending', 'msgqueue', 'nice', 'rtprio', 'chroot']

    pam_types = [ 'soft', 'hard', '-' ]

    limits_conf = '/etc/security/limits.conf'

    module = AnsibleModule(
        # not checking because of daisy chain to file module
        argument_spec = dict(
            domain            = dict(required=True, type='str'),
            limit_type        = dict(required=True, type='str', choices=pam_types),
            limit_item        = dict(required=True, type='str', choices=pam_items),
            value             = dict(required=True, type='str'),
            use_max           = dict(default=False, type='bool'),
            use_min           = dict(default=False, type='bool'),
            backup            = dict(default=False, type='bool'),
            dest              = dict(default=limits_conf, type='str'),
            comment           = dict(required=False, default='', type='str')
        )
    )

    domain      =       module.params['domain']
    limit_type  =       module.params['limit_type']
    limit_item  =       module.params['limit_item']
    value       =       module.params['value']
    use_max     =       module.params['use_max']
    use_min     =       module.params['use_min']
    backup      =       module.params['backup']
    limits_conf =       module.params['dest']
    new_comment =       module.params['comment']

    changed = False

    if os.path.isfile(limits_conf):
        if not os.access(limits_conf, os.W_OK):
            module.fail_json(msg="%s is not writable. Use sudo" % (limits_conf) )
    else:
        module.fail_json(msg="%s is not visible (check presence, access rights, use sudo)" % (limits_conf) )

    if use_max and use_min:
        module.fail_json(msg="Cannot use use_min and use_max at the same time." )

    if not (value in ['unlimited', 'infinity', '-1'] or value.isdigit()):
        module.fail_json(msg="Argument 'value' can be one of 'unlimited', 'infinity', '-1' or positive number. Refer to manual pages for more details.")

    # Backup
    if backup:
        backup_file = module.backup_local(limits_conf)

    space_pattern = re.compile(r'\s+')

    message = ''
    f = open (limits_conf, 'rb')
    # Tempfile
    nf = tempfile.NamedTemporaryFile(mode='w+')

    found = False
    new_value = value

    for line in f:
        line = to_native(line, errors='surrogate_or_strict')
        if line.startswith('#'):
            nf.write(line)
            continue

        newline = re.sub(space_pattern, ' ', line).strip()
        if not newline:
            nf.write(line)
            continue

        # Remove comment in line
        newline = newline.split('#',1)[0]
        try:
            old_comment = line.split('#',1)[1]
        except:
            old_comment = ''

        newline = newline.rstrip()

        if not new_comment:
            new_comment = old_comment

        line_fields = newline.split(' ')

        if len(line_fields) != 4:
            nf.write(line)
            continue

        line_domain     = line_fields[0]
        line_type       = line_fields[1]
        line_item       = line_fields[2]
        actual_value    = line_fields[3]

        if not (actual_value in ['unlimited', 'infinity', '-1'] or actual_value.isdigit()):
            module.fail_json(msg="Invalid configuration of '%s'. Current value of %s is unsupported." % (limits_conf, line_item))

        # Found the line
        if line_domain == domain and line_type == limit_type and line_item == limit_item:
            found = True
            if value == actual_value:
                message = line
                nf.write(line)
                continue

            actual_value_unlimited = actual_value in ['unlimited', 'infinity', '-1']
            value_unlimited = value in ['unlimited', 'infinity', '-1']

            if use_max:
                if value.isdigit() and actual_value.isdigit():
                    new_value = str(max(int(value), int(actual_value)))
                elif actual_value_unlimited:
                    new_value = actual_value
                else:
                    new_value = value

            if use_min:
                if value.isdigit() and actual_value.isdigit():
                    new_value = str(min(int(value), int(actual_value)))
                elif value_unlimited:
                    new_value = actual_value
                else:
                    new_value = value

            # Change line only if value has changed
            if new_value != actual_value:
                changed = True
                if new_comment:
                    new_comment = "\t#" + new_comment
                new_limit = domain + "\t" + limit_type + "\t" + limit_item + "\t" + new_value + new_comment + "\n"
                message = new_limit
                nf.write(new_limit)
            else:
                message = line
                nf.write(line)
        else:
            nf.write(line)

    if not found:
        changed = True
        if new_comment:
            new_comment = "\t#"+new_comment
        new_limit = domain + "\t" + limit_type + "\t" + limit_item + "\t" + new_value + new_comment + "\n"
        message = new_limit
        nf.write(new_limit)

    f.close()
    nf.flush()

    # Copy tempfile to newfile
    module.atomic_move(nf.name, f.name)

    try:
        nf.close()
    except:
        pass

    res_args = dict(
        changed = changed, msg = message
    )

    if backup:
        res_args['backup_file'] = backup_file

    module.exit_json(**res_args)

if __name__ == '__main__':
    main()
<commit_msg>pam_limits creates custom file if not exist (#26989) (#28616)<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2014, Sebastien Rohaut <sebastien.rohaut@gmail.com>
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function
__metaclass__ = type

ANSIBLE_METADATA = {'metadata_version': '1.1',
                    'status': ['preview'],
                    'supported_by': 'community'}

DOCUMENTATION = '''
---
module: pam_limits
version_added: "2.0"
author:
    - "Sebastien Rohaut (@usawa)"
short_description: Modify Linux PAM limits
description:
     - The C(pam_limits) module modify PAM limits, default in /etc/security/limits.conf.
       For the full documentation, see man limits.conf(5).
options:
  domain:
    description:
      - A username, @groupname, wildcard, uid/gid range.
    required: true
  limit_type:
    description:
      - Limit type, see C(man limits) for an explanation
    required: true
    choices: [ "hard", "soft", "-" ]
  limit_item:
    description:
      - The limit to be set
    required: true
    choices:
        - "core"
        - "data"
        - "fsize"
        - "memlock"
        - "nofile"
        - "rss"
        - "stack"
        - "cpu"
        - "nproc"
        - "as"
        - "maxlogins"
        - "maxsyslogins"
        - "priority"
        - "locks"
        - "sigpending"
        - "msgqueue"
        - "nice"
        - "rtprio"
        - "chroot"
  value:
    description:
      - The value of the limit.
    required: true
  backup:
    description:
      - Create a backup file including the timestamp information so you can get
        the original file back if you somehow clobbered it incorrectly.
    required: false
    choices: [ "yes", "no" ]
    default: "no"
  use_min:
    description:
      - If set to C(yes), the minimal value will be used or conserved.
        If the specified value is inferior to the value in the file, file content is replaced with the new value,
        else content is not modified.
    required: false
    choices: [ "yes", "no" ]
    default: "no"
  use_max:
    description:
      - If set to C(yes), the maximal value will be used or conserved.
        If the specified value is superior to the value in the file, file content is replaced with the new value,
        else content is not modified.
    required: false
    choices: [ "yes", "no" ]
    default: "no"
  dest:
    description:
      - Modify the limits.conf path.
    required: false
    default: "/etc/security/limits.conf"
  comment:
    description:
      - Comment associated with the limit.
    required: false
    default: ''
notes:
  - If dest file doesn't exists, it is created.
'''

EXAMPLES = '''
# Add or modify nofile soft limit for the user joe
- pam_limits:
    domain: joe
    limit_type: soft
    limit_item: nofile
    value: 64000

# Add or modify fsize hard limit for the user smith. Keep or set the maximal value.
- pam_limits:
    domain: smith
    limit_type: hard
    limit_item: fsize
    value: 1000000
    use_max: yes

# Add or modify memlock, both soft and hard, limit for the user james with a comment.
- pam_limits:
    domain: james
    limit_type: '-'
    limit_item: memlock
    value: unlimited
    comment: unlimited memory lock for james
'''

import os
import os.path
import tempfile
import re

from ansible.module_utils.basic import AnsibleModule
from ansible.module_utils._text import to_native

def main():

    pam_items = ['core', 'data', 'fsize', 'memlock', 'nofile', 'rss', 'stack', 'cpu', 'nproc', 'as', 'maxlogins', 'maxsyslogins', 'priority', 'locks',
                 'sigpending', 'msgqueue', 'nice', 'rtprio', 'chroot']

    pam_types = [ 'soft', 'hard', '-' ]

    limits_conf = '/etc/security/limits.conf'

    module = AnsibleModule(
        # not checking because of daisy chain to file module
        argument_spec = dict(
            domain            = dict(required=True, type='str'),
            limit_type        = dict(required=True, type='str', choices=pam_types),
            limit_item        = dict(required=True, type='str', choices=pam_items),
            value             = dict(required=True, type='str'),
            use_max           = dict(default=False, type='bool'),
            use_min           = dict(default=False, type='bool'),
            backup            = dict(default=False, type='bool'),
            dest              = dict(default=limits_conf, type='str'),
            comment           = dict(required=False, default='', type='str')
        )
    )

    domain      =       module.params['domain']
    limit_type  =       module.params['limit_type']
    limit_item  =       module.params['limit_item']
    value       =       module.params['value']
    use_max     =       module.params['use_max']
    use_min     =       module.params['use_min']
    backup      =       module.params['backup']
    limits_conf =       module.params['dest']
    new_comment =       module.params['comment']

    changed = False

    if os.path.isfile(limits_conf):
        if not os.access(limits_conf, os.W_OK):
            module.fail_json(msg="%s is not writable. Use sudo" % limits_conf)
    else:
        limits_conf_dir = os.path.dirname(limits_conf)
        if os.path.isdir(limits_conf_dir) and os.access(limits_conf_dir, os.W_OK):
            open(limits_conf, 'a').close()
            changed = True
        else:
            module.fail_json(msg="directory %s is not writable (check presence, access rights, use sudo)" % limits_conf_dir)

    if use_max and use_min:
        module.fail_json(msg="Cannot use use_min and use_max at the same time.")

    if not (value in ['unlimited', 'infinity', '-1'] or value.isdigit()):
        module.fail_json(msg="Argument 'value' can be one of 'unlimited', 'infinity', '-1' or positive number. Refer to manual pages for more details.")

    # Backup
    if backup:
        backup_file = module.backup_local(limits_conf)

    space_pattern = re.compile(r'\s+')

    message = ''
    f = open (limits_conf, 'rb')
    # Tempfile
    nf = tempfile.NamedTemporaryFile(mode='w+')

    found = False
    new_value = value

    for line in f:
        line = to_native(line, errors='surrogate_or_strict')
        if line.startswith('#'):
            nf.write(line)
            continue

        newline = re.sub(space_pattern, ' ', line).strip()
        if not newline:
            nf.write(line)
            continue

        # Remove comment in line
        newline = newline.split('#',1)[0]
        try:
            old_comment = line.split('#',1)[1]
        except:
            old_comment = ''

        newline = newline.rstrip()

        if not new_comment:
            new_comment = old_comment

        line_fields = newline.split(' ')

        if len(line_fields) != 4:
            nf.write(line)
            continue

        line_domain     = line_fields[0]
        line_type       = line_fields[1]
        line_item       = line_fields[2]
        actual_value    = line_fields[3]

        if not (actual_value in ['unlimited', 'infinity', '-1'] or actual_value.isdigit()):
            module.fail_json(msg="Invalid configuration of '%s'. Current value of %s is unsupported." % (limits_conf, line_item))

        # Found the line
        if line_domain == domain and line_type == limit_type and line_item == limit_item:
            found = True
            if value == actual_value:
                message = line
                nf.write(line)
                continue

            actual_value_unlimited = actual_value in ['unlimited', 'infinity', '-1']
            value_unlimited = value in ['unlimited', 'infinity', '-1']

            if use_max:
                if value.isdigit() and actual_value.isdigit():
                    new_value = str(max(int(value), int(actual_value)))
                elif actual_value_unlimited:
                    new_value = actual_value
                else:
                    new_value = value

            if use_min:
                if value.isdigit() and actual_value.isdigit():
                    new_value = str(min(int(value), int(actual_value)))
                elif value_unlimited:
                    new_value = actual_value
                else:
                    new_value = value

            # Change line only if value has changed
            if new_value != actual_value:
                changed = True
                if new_comment:
                    new_comment = "\t#" + new_comment
                new_limit = domain + "\t" + limit_type + "\t" + limit_item + "\t" + new_value + new_comment + "\n"
                message = new_limit
                nf.write(new_limit)
            else:
                message = line
                nf.write(line)
        else:
            nf.write(line)

    if not found:
        changed = True
        if new_comment:
            new_comment = "\t#"+new_comment
        new_limit = domain + "\t" + limit_type + "\t" + limit_item + "\t" + new_value + new_comment + "\n"
        message = new_limit
        nf.write(new_limit)

    f.close()
    nf.flush()

    # Copy tempfile to newfile
    module.atomic_move(nf.name, f.name)

    try:
        nf.close()
    except:
        pass

    res_args = dict(
        changed = changed, msg = message
    )

    if backup:
        res_args['backup_file'] = backup_file

    module.exit_json(**res_args)

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


<commit_before>"""Classes for managing bit torrent data"""
from functools import reduce
from collections import namedtuple
from pizza_utils.listutils import split, chunk
from pizza_utils.bitfield import Bitfield
from typing import Optional, Tuple
from . import bencoding, peer, tracker
import aiohttp.client_exceptions
import hashlib
import logging
import math
import operator
import random

BLOCK_REQUEST_SIZE = 2**15  # Bytes

TorrentFile = namedtuple("TorrentFile", ["path", "filename", "size", "piece_pointer"])
TorrentFile.__doc__ = """A file indicated by the torrent metadata"""
TorrentFile.path.__doc__ = """The path (relative to the working directory) or the file."""
TorrentFile.filename.__doc__ = """The name of the file on disk."""
TorrentFile.size.__doc__ = """The size of the file in bytes."""
TorrentFile.piece_pointer.__doc__ = "The index of the first piece that holds data for this file."

def calculate_rarity(peer_list: list, piece_count: int):
    """Maps rarity to a list of piece indexes"""
    pieces = {}

    for i in range(0, piece_count):
        rarity = 0
        for p in peer_list:
            if p.piecefield[i] > 0:
                rarity = rarity + 1

        if rarity in pieces:
            pieces[rarity].append(i)
        else:
            pieces[rarity] = [i]

    return pieces

class InvalidTorrentError(Exception):
    """Exception for errors involving malformed torrent files.

    Attributes:
        filename -- the file that caused this error
        message -- explaination of error
    """

    def __init__(self, filename, message):
        self.filename = filename
        self.message = message

class TorrentMeta:
    """
    Helper class for dealing with torrent metadata.

    """
    def __init__(self, filename: str):
        """
        Creates a new TorrentMeta object
        :param filename: path to the torrent file
        """
        self.filename = filename

        with open(filename, mode="rb") as f:
            self._meta = bencoding.decode(f.read())[0]

        self._init_announce_urls()
        self._init_pieces()
        self._init_pieces()

    def _init_announce_urls(self):
        self.announce = []

        if b"announce-list" in self._meta:
            for announce_list in self._meta[b"announce-list"]:
                self.announce.append([url.decode() for url in announce_list])
        elif b"announce" in self._meta:
            self.announce.append([self._meta[b"announce"].decode()])
        else:
            raise InvalidTorrentError(self.filename, "File does not contain an announce.")

    def _init_pieces(self):
        info: dict = self._meta[b"info"]

        if b"pieces" in info:
            self.pieces = tuple(chunk(info[b"pieces"], 20))
        else:
            raise InvalidTorrentError(self.filename, "File does not contains pieces.")

    def _init_files(self):
        pass

    @property
    def piece_length(self) -> int:
        """
        :return: The size of the pieces in bytes
        """
        info: dict = self._meta[b"info"] # Do this to appease Pycharm
        return info[b"piece length"]

    @property
    def creation_date(self) -> int:
        """
        :return: The creation time of the torrent in UNIX epoch format.
        """
        return self._meta[b"creation date"]

    @property
    def comment(self) -> str:
        return self._meta[b"comment"].decode()

    @property
    def created_by(self) -> str:
        return self._meta[b"created by"].decode()

    @property
    def encoding(self) -> str:
        return self._meta[b"encoding"].decode()

class Piece():
    """
    Class that represents a piece, and provides a buffer for downloading.

    Attributes:
        downloaded      The piece's data has been downloaded, but not saved
        verified        The piece's data has been downloaded and verified, but not saved
        saved           The piece has been downloaded, verified, and saved.
    """

    def __init__(self, piece_hash: bytes, piece_size: Optional[int]):
        self.piece_hash = piece_hash
        self.piece_size = piece_size

        self.recycle()

    def __repr__(self):
        if self.downloaded:
            return "Piece {hash}: {buffer}".format(self.piece_hash, self.buffer)
        elif self.verified:
            return "Piece (v) {hask}: {buffer}".format(self.piece_hash, self.buffer)
        elif self.downloaded:
            return "Piece (v,s) {hask}: {buffer}".format(self.piece_hash, self.buffer)
        else:
            return "Piece {hash}".format(self.piece_hash)

    def next_offset(self) -> int:
        return self.piece_size - len(self.data)

    def recycle(self):
        self.buffer = bytearray()
        self.downloaded = False
        self.verified = False
        self.saved = False

    async def download(self, offset: int, data: bytes):
        self.buffer[offset:offset + len(data)] = data

        if len(self.buffer) == self.piece_size:
            self.downloaded = True

    async def verify(self) -> bool:
        if hashlib.sha1(self.buffer).digest() == self.piece_hash:
            self.verified = True
            return True
        else:
            return False

    async def save(self, file: TorrentFile, offset: int = 0):
        byte_offset = offset * self.piece_size

        with open(file.path + file.name, mode="w+b") as fh:
            fh.seek(byte_offset)
            fh.write(self.buffer)

        # Clean up and free memory
        del self.buffer
        self.buffer = bytearray()

        self.saved = True

class Torrent:
    """
    Class for storing Torrent state.

    Attributes:
        data            The TorrentData of the Torrent
        swarm           A list of all the peers.
        peers           A list of the peers the client is connected too
        pieces          A list of the pieces
        file_indexes    A map of piece indexes and files
    """

    def __init__(self, filename: str):
        self.data = TorrentMeta(filename)
        self.swarm = []
        self.peers = []

        self.total_uploaded = 0
        self.total_downloaded = 0

        self.pieces = tuple(Piece(ph, self.data["info.piece length"]) for ph in self.data.pieces)
        self.downloading = []
        # File indexes are basically pointers to items in the pieces list
        # They're indicators of which pieces correspond to which files
        self.file_indexes = self._calculate_file_indexes()

        # Unique key to identify the torrent
        self.key = bytes(random.sample(range(0, 256), 8))

        self._logger = logging.getLogger("bridge.torrent." + self.data.name)
        self._announce_time = []
        self._celebrate = False

    def __str__(self):
        return "Torrent(name: {}, swarm: {}, peers: {})".format(self.data.name, len(self.swarm), len(self.peers))

    @property
    def downloaded(self) -> int:
        rv = 0
        for p in self.pieces:
            if p.saved:
                rv = rv + p.piece_size
        return rv

    @property
    def left(self) -> int:
        rv = 0
        for p in self.pieces:
            if not p.verified:
                rv = rv + p.piece_size
        return rv

    def _calculate_file_indexes(self):
        rv = {}
        next_index = 0

        for file in self.data.files:
            rv[next_index] = file
            next_index = next_index + math.ceil(file.size / self.data["info.piece length"])

    @property
    def bitfield(self) -> Bitfield:
        rv = Bitfield()
        for p in self.pieces:
            if p.saved:
                rv[0] = 1
            else:
                rv = rv << 1

        return rv

    def get_piece_file(self, i: int) -> Tuple[int, TorrentFile]:
        for index, file in reversed(self.file_indexes):
            if i > index:
                return (index, file)

    def ask_for_block(self, remote_peer) -> peer.PeerMessage:
        for i in self.downloading:
            if remote_peer.piecefield[i] > 0:
                offset = len(self.pieces[i].buffer)
                self._logger.debug("Asking for piece {} at offset {} from {}".format(i, offset, remote_peer))
                return peer.RequestPeerMessage(i, offset, BLOCK_REQUEST_SIZE)

        rp = calculate_rarity(self.swarm, len(self.pieces))
        rarity = sorted(rp.keys(), reverse=True)

        for r in rarity:
            for i in rp[r]:
                if remote_peer.piecefield[i] > 0 and i not in self.downloading and not self.pieces[i].verified:
                    self._logger.debug("Asking for piece {} from {}".format(i, remote_peer))
                    return peer.RequestPeerMessage(i, 0, BLOCK_REQUEST_SIZE)

    async def recieve_block(self, piece_index: int, offset: int, data: bytes):
        """Handler for recieving a block of data"""
        p: Piece = self.pieces[piece_index]
        self.downloading.append(piece_index)
        self.total_downloaded = self.total_downloaded + len(data)

        await p.download(offset, data)

        if p.downloaded:
            self._logger.debug("Downloaded {}".format(p))

            if await p.verify():
                # The piece has been fully downloaded and verified, save to disk
                self._logger.debug("Verified {}".format(p))

                f_index, f = self.get_piece_file(piece_index)
                piece_offset = piece_index - f_index

                self._logger.debug("Saving {} to {} offset {}".format(p, f, piece_offset))

                try:
                    p.save(f, piece_offset)
                    self._logger.info("Successfully downloaded piece {}.".format_map(piece_index))

                    # Because why not
                    if self._celebrate is False:
                        print("Congrats! You just downloaded your first piece! :)")
                        print("Here it is:")
                        print(p.buffer)

                        self._celebrate = True
                except OSError:
                    self._logger.warn("Error downloading piece {}. Discarding.")
                    p.recycle()
                finally:
                    self.downloading.remove(piece_index)
            else:
                self._logger.debug("Piece {} failed to verify. Recycling.".format(p))
                p.recycle()

    async def announce(self, port: int, peer_id: bytes):
        self._logger.info("Beginning anounce...")

        # TODO: Add support for backups
        # Normally, you go through the first sub-list of URLs and then move to the second sub-list only if all the announces
        # fail in the first list. Then, the sub-lists are rearragned so that the first successful trackers are first in
        # their sub-lists. http://bittorrent.org/beps/bep_0012.html
        for announce_url in self.data.announce[0]:
            self._logger.info("Announcing on " + announce_url)

            peer_count_request = max(peer.NEW_CONNECTION_LIMIT - len(self.peers), 0)
            self._logger.info("Requesting {} peers".format(peer_count_request))

            try:
                response = await tracker.announce_tracker(self, peer_id, announce_url, port, numwant=peer_count_request)

                if not response.sucessful:
                    self._logger.warning("Announce on {} not successful. {}".format(announce_url, str(response)))
                    continue

                self._logger.info("Announce successful.")

                if peer_count_request > 0:
                    # Dump all the peers into the swarm
                    new_peers = [p for p in response.peers if p not in self.swarm]
                    self.swarm.extend(new_peers)
                    self._logger.info("Adding {} new peers into the swarm. (Now {})".format(len(new_peers), len(self.swarm)))
            except aiohttp.client_exceptions.ClientConnectionError:
                self._logger.error("Failed to connect to tracker {}".format(announce_url))
                continue
<commit_msg>Added files into TorrentMeta<commit_after>"""Classes for managing bit torrent data"""
from functools import reduce
from collections import namedtuple
from pizza_utils.listutils import split, chunk
from pizza_utils.bitfield import Bitfield
from typing import Optional, Tuple
from . import bencoding, peer, tracker
import aiohttp.client_exceptions
import hashlib
import logging
import math
import operator
import random

BLOCK_REQUEST_SIZE = 2**15  # Bytes

TorrentFile = namedtuple("TorrentFile", ["path", "filename", "size", "piece_pointer"])
TorrentFile.__doc__ = """A file indicated by the torrent metadata"""
TorrentFile.path.__doc__ = """The path (relative to the working directory) or the file."""
TorrentFile.filename.__doc__ = """The name of the file on disk."""
TorrentFile.size.__doc__ = """The size of the file in bytes."""
TorrentFile.piece_pointer.__doc__ = "The index of the first piece that holds data for this file."

def calculate_rarity(peer_list: list, piece_count: int):
    """Maps rarity to a list of piece indexes"""
    pieces = {}

    for i in range(0, piece_count):
        rarity = 0
        for p in peer_list:
            if p.piecefield[i] > 0:
                rarity = rarity + 1

        if rarity in pieces:
            pieces[rarity].append(i)
        else:
            pieces[rarity] = [i]

    return pieces

class InvalidTorrentError(Exception):
    """Exception for errors involving malformed torrent files.

    Attributes:
        filename -- the file that caused this error
        message -- explaination of error
    """

    def __init__(self, filename, message):
        self.filename = filename
        self.message = message

class TorrentMeta:
    """
    Helper class for dealing with torrent metadata.

    """
    def __init__(self, filename: str):
        """
        Creates a new TorrentMeta object
        :param filename: path to the torrent file
        """
        self.filename = filename

        with open(filename, mode="rb") as f:
            self._meta = bencoding.decode(f.read())[0]

        self._init_announce_urls()
        self._init_pieces()
        self._init_pieces()

    def _init_announce_urls(self):
        self.announce = []

        if b"announce-list" in self._meta:
            for announce_list in self._meta[b"announce-list"]:
                self.announce.append([url.decode() for url in announce_list])
        elif b"announce" in self._meta:
            self.announce.append([self._meta[b"announce"].decode()])
        else:
            raise InvalidTorrentError(self.filename, "File does not contain an announce.")

    def _init_pieces(self):
        info: dict = self._meta[b"info"]

        if b"pieces" in info:
            self.pieces = tuple(chunk(info[b"pieces"], 20))
        else:
            raise InvalidTorrentError(self.filename, "File does not contain pieces.")

    def _init_files(self):
        self.files = []
        info: dict = self._meta[b"info"]

        if b"files" in info:
            # Multi-file mode
            piece_pointer = 0

            for item in info[b"files"]:
                path = [s.decode() for s in item["path"]]
                size = item[b"length"]

                self.files.append(TorrentFile("/".join(path[:-1]), path[-1], size, piece_pointer))

                piece_pointer +=  math.floor(size / self.piece_length)
        elif b"name" in info and b"length" in info:
            self.files =[TorrentFile("", info[b"name"].decode(), info[b"length"], 0)]
        else:
            raise InvalidTorrentError(self.filename, "File does not contain files.")

    @property
    def piece_length(self) -> int:
        """
        :return: The size of the pieces in bytes
        """
        info: dict = self._meta[b"info"] # Do this to appease Pycharm
        return info[b"piece length"]

    @property
    def creation_date(self) -> int:
        """
        :return: The creation time of the torrent in UNIX epoch format.
        """
        return self._meta[b"creation date"]

    @property
    def comment(self) -> str:
        return self._meta[b"comment"].decode()

    @property
    def created_by(self) -> str:
        return self._meta[b"created by"].decode()

    @property
    def encoding(self) -> str:
        return self._meta[b"encoding"].decode()

class Piece():
    """
    Class that represents a piece, and provides a buffer for downloading.

    Attributes:
        downloaded      The piece's data has been downloaded, but not saved
        verified        The piece's data has been downloaded and verified, but not saved
        saved           The piece has been downloaded, verified, and saved.
    """

    def __init__(self, piece_hash: bytes, piece_size: Optional[int]):
        self.piece_hash = piece_hash
        self.piece_size = piece_size

        self.recycle()

    def __repr__(self):
        if self.downloaded:
            return "Piece {hash}: {buffer}".format(self.piece_hash, self.buffer)
        elif self.verified:
            return "Piece (v) {hask}: {buffer}".format(self.piece_hash, self.buffer)
        elif self.downloaded:
            return "Piece (v,s) {hask}: {buffer}".format(self.piece_hash, self.buffer)
        else:
            return "Piece {hash}".format(self.piece_hash)

    def next_offset(self) -> int:
        return self.piece_size - len(self.data)

    def recycle(self):
        self.buffer = bytearray()
        self.downloaded = False
        self.verified = False
        self.saved = False

    async def download(self, offset: int, data: bytes):
        self.buffer[offset:offset + len(data)] = data

        if len(self.buffer) == self.piece_size:
            self.downloaded = True

    async def verify(self) -> bool:
        if hashlib.sha1(self.buffer).digest() == self.piece_hash:
            self.verified = True
            return True
        else:
            return False

    async def save(self, file: TorrentFile, offset: int = 0):
        byte_offset = offset * self.piece_size

        with open(file.path + file.name, mode="w+b") as fh:
            fh.seek(byte_offset)
            fh.write(self.buffer)

        # Clean up and free memory
        del self.buffer
        self.buffer = bytearray()

        self.saved = True

class Torrent:
    """
    Class for storing Torrent state.

    Attributes:
        data            The TorrentData of the Torrent
        swarm           A list of all the peers.
        peers           A list of the peers the client is connected too
        pieces          A list of the pieces
        file_indexes    A map of piece indexes and files
    """

    def __init__(self, filename: str):
        self.data = TorrentMeta(filename)
        self.swarm = []
        self.peers = []

        self.total_uploaded = 0
        self.total_downloaded = 0

        self.pieces = tuple(Piece(ph, self.data["info.piece length"]) for ph in self.data.pieces)
        self.downloading = []
        # File indexes are basically pointers to items in the pieces list
        # They're indicators of which pieces correspond to which files
        self.file_indexes = self._calculate_file_indexes()

        # Unique key to identify the torrent
        self.key = bytes(random.sample(range(0, 256), 8))

        self._logger = logging.getLogger("bridge.torrent." + self.data.name)
        self._announce_time = []
        self._celebrate = False

    def __str__(self):
        return "Torrent(name: {}, swarm: {}, peers: {})".format(self.data.name, len(self.swarm), len(self.peers))

    @property
    def downloaded(self) -> int:
        rv = 0
        for p in self.pieces:
            if p.saved:
                rv = rv + p.piece_size
        return rv

    @property
    def left(self) -> int:
        rv = 0
        for p in self.pieces:
            if not p.verified:
                rv = rv + p.piece_size
        return rv

    def _calculate_file_indexes(self):
        rv = {}
        next_index = 0

        for file in self.data.files:
            rv[next_index] = file
            next_index = next_index + math.ceil(file.size / self.data["info.piece length"])

    @property
    def bitfield(self) -> Bitfield:
        rv = Bitfield()
        for p in self.pieces:
            if p.saved:
                rv[0] = 1
            else:
                rv = rv << 1

        return rv

    def get_piece_file(self, i: int) -> Tuple[int, TorrentFile]:
        for index, file in reversed(self.file_indexes):
            if i > index:
                return (index, file)

    def ask_for_block(self, remote_peer) -> peer.PeerMessage:
        for i in self.downloading:
            if remote_peer.piecefield[i] > 0:
                offset = len(self.pieces[i].buffer)
                self._logger.debug("Asking for piece {} at offset {} from {}".format(i, offset, remote_peer))
                return peer.RequestPeerMessage(i, offset, BLOCK_REQUEST_SIZE)

        rp = calculate_rarity(self.swarm, len(self.pieces))
        rarity = sorted(rp.keys(), reverse=True)

        for r in rarity:
            for i in rp[r]:
                if remote_peer.piecefield[i] > 0 and i not in self.downloading and not self.pieces[i].verified:
                    self._logger.debug("Asking for piece {} from {}".format(i, remote_peer))
                    return peer.RequestPeerMessage(i, 0, BLOCK_REQUEST_SIZE)

    async def recieve_block(self, piece_index: int, offset: int, data: bytes):
        """Handler for recieving a block of data"""
        p: Piece = self.pieces[piece_index]
        self.downloading.append(piece_index)
        self.total_downloaded = self.total_downloaded + len(data)

        await p.download(offset, data)

        if p.downloaded:
            self._logger.debug("Downloaded {}".format(p))

            if await p.verify():
                # The piece has been fully downloaded and verified, save to disk
                self._logger.debug("Verified {}".format(p))

                f_index, f = self.get_piece_file(piece_index)
                piece_offset = piece_index - f_index

                self._logger.debug("Saving {} to {} offset {}".format(p, f, piece_offset))

                try:
                    p.save(f, piece_offset)
                    self._logger.info("Successfully downloaded piece {}.".format_map(piece_index))

                    # Because why not
                    if self._celebrate is False:
                        print("Congrats! You just downloaded your first piece! :)")
                        print("Here it is:")
                        print(p.buffer)

                        self._celebrate = True
                except OSError:
                    self._logger.warn("Error downloading piece {}. Discarding.")
                    p.recycle()
                finally:
                    self.downloading.remove(piece_index)
            else:
                self._logger.debug("Piece {} failed to verify. Recycling.".format(p))
                p.recycle()

    async def announce(self, port: int, peer_id: bytes):
        self._logger.info("Beginning anounce...")

        # TODO: Add support for backups
        # Normally, you go through the first sub-list of URLs and then move to the second sub-list only if all the announces
        # fail in the first list. Then, the sub-lists are rearragned so that the first successful trackers are first in
        # their sub-lists. http://bittorrent.org/beps/bep_0012.html
        for announce_url in self.data.announce[0]:
            self._logger.info("Announcing on " + announce_url)

            peer_count_request = max(peer.NEW_CONNECTION_LIMIT - len(self.peers), 0)
            self._logger.info("Requesting {} peers".format(peer_count_request))

            try:
                response = await tracker.announce_tracker(self, peer_id, announce_url, port, numwant=peer_count_request)

                if not response.sucessful:
                    self._logger.warning("Announce on {} not successful. {}".format(announce_url, str(response)))
                    continue

                self._logger.info("Announce successful.")

                if peer_count_request > 0:
                    # Dump all the peers into the swarm
                    new_peers = [p for p in response.peers if p not in self.swarm]
                    self.swarm.extend(new_peers)
                    self._logger.info("Adding {} new peers into the swarm. (Now {})".format(len(new_peers), len(self.swarm)))
            except aiohttp.client_exceptions.ClientConnectionError:
                self._logger.error("Failed to connect to tracker {}".format(announce_url))
                continue
<|endoftext|>


<commit_before>

import os, os.path, string, tempfile, traceback,gtk,gtk.glade,sys

#sys.path.insert(0,"/usr/lib/pygtk1/lib/python"+sys.version[:3]+"/site-packages/gtk-1.2")
#sys.path.insert(0,os.path.abspath(os.path.dirname(sys.argv[0]))+"/../lib/python"+sys.version[:3]+"/site-packages/gtk-1.2")
#sys.path.insert(0, PYGTK1PLACEMENT  +"/../lib/python"+sys.version[:3]+"/site-packages/gtk-1.2")

def lrstrip(dasstring):
    dasstring=string.lstrip(dasstring)
    dasstring=string.rstrip(dasstring)
    return dasstring

class Config:
    def __init__(self,xmltree,conffilename):
        self.xmltree=xmltree
        self.vars={}
        self.conffilename=conffilename
        self.readConfFile(conffilename)
        self.fileselectors={}

    def setTempFileName(self,tempfilename):
        self.tempfilename=tempfilename

    def readConfFile(self,filename):
        try:
            file=open(filename,"r")
        except:
            return
        while 1:
            line=""
            while line=="" or line=="\n" or line[0:1]=="#":
                line=file.readline()
                if line=="":
                    file.close()
                    return
                line=lrstrip(line)

            splitline=string.split(line,"=",1)
            splitline[0]=lrstrip(splitline[0])
            splitline[1]=lrstrip(splitline[1])

            if isinstance(self.xmltree.get_widget(splitline[0]),get.Entry):
                self.vars[splitline[0]]=splitline[1]
            else:
                self.vars[splitline[0]]=int(splitline[1])

    def writeConfFile(self,filename):
        try:
            file=open(filename,"w")
        except:
            print "Could not write file "+filename

        varnames=self.vars.keys()
        varnames.sort()
        for varname in varnames:
            if isinstance(self.xmltree.get_widget(varname), gtk.Entry):
                file.write(varname+"="+self.vars[varname]+"\n")
            else:
                file.write("%s=%d\n" % (varname,self.vars[varname]))

        try:
            file.close()
        except:
            print "Could not write file "+filename

    def getVar(self,confname):
        if self.vars.has_key(confname):
            return self.vars[confname]
        widg=self.xmltree.get_widget(confname)
        if widg!=None:
            if isinstance(widg,gtk.Entry):
                return widg.get_text()
            else:
                return widg.props.active

    def generalHandler(self,widg):
        name=widg.get_name()

        if isinstance(widg, gtk.Button):
            if name=="Ok":
                self.writeConfFile(self.tempfilename)
                gtk.main_quit()
                return
            if name=="Cancel":
                gtk.main_quit()
                return
            if name=="Save":
                self.writeConfFile(self.conffilename)
                self.writeConfFile(self.tempfilename)
                if self.getVar("close_window_when_pressing_save")==1:
                    gtk.main_quit()
                return

            if name[:7]=="browse+":
                widgnamestrip=string.split(name,"+",1)
                widgname=widgnamestrip[1]
                self.fileselectors[widgname]=gtk.glade.XML(
                    os.path.abspath(os.path.dirname(sys.argv[0]))+"/../etc/ceresconfig.glade",widgname
                    )
                self.fileselectors[widgname].signal_autoconnect(handle_dict)
                widg=self.fileselectors[widgname].get_widget(widgname)
                textentry=self.xmltree.get_widget(widgname[3:])
                widg.set_filename(textentry.get_text())
            else:
                widgnamestrip=string.split(name,"-",1)
                widgname="fs_"+widgnamestrip[1]
                widg=self.fileselectors[widgname].get_widget(widgname)
                if widgnamestrip[0]=="fsok":
                    filename=widg.get_filename()
                    if widgname!="fs_frequency_file":
                        filename=os.path.dirname(filename)+"/"
                    textentry=self.xmltree.get_widget(widgnamestrip[1])
                    textentry.set_text(filename)
                    self.vars[widgnamestrip[1]]=filename
                widg.destroy()
                del self.fileselectors[widgname]
        else:
            if isinstance(widg,gtk.Entry):
                widgtext=widg.get_text()
                if name[:-1]=="temporary_path_":
                    if widgtext[-1]!="/":
                        widgtext=widgtext+"/"
                        widg.set_text(widgtext)
                self.vars[name]=widgtext
            else:
                self.vars[name]=widg.props.active

    def destroy(self):
        gtk.main_quit()

    def show(self):
        self.xmltree.signal_autoconnect(handle_dict)
        varnames=self.vars.keys()
        for varname in varnames:
            widget=self.xmltree.get_widget(varname)
            if isinstance(widget,  gtk.GtkEntry):
                widget.set_text(self.vars[varname])
            else:
                widget.set_active(self.vars[varname])

        gtk.main()

def getVar(confname):
    try:
        return config.getVar(confname)
    except:
        traceback.print_exc(file=sys.stderr)

def generalHandler(obj):
    config.generalHandler(obj)

def destroy(obj):
    config.destroy()

def showConfig():
    try:
        tempfilename=tempfile.mktemp()
        config.writeConfFile(tempfilename)
        os.system(
            os.path.abspath(os.path.dirname(sys.argv[0]))
            + "/ceresconfig.py %s %s" % (config.conffilename,tempfilename)
            )
        config.readConfFile(config.conffilename)
        config.readConfFile(tempfilename)
        os.unlink(tempfilename)
    except:
        traceback.print_exc(file=sys.stderr)

handle_dict={
    "general":generalHandler,
    "on_Settings_destroy":destroy
    }

def ceresConfigStart(filename):
    global config
    a=sys.argv[0]
    b=os.path.dirname(a)
    c=os.path.abspath(b)
    d=c+"/../etc/ceresconfig.glade"
    config=Config(
        gtk.glade.XML(d,"Settings"),
        filename
        )

if __name__=="__main__":
    global config

    if len(sys.argv)>1 and (sys.argv[1]=="--help" or sys.argv[1]=="-h" or sys.argv[1]=="-help"):
        print
        print "Usage: %s [savefilename] [okfilename]" % sys.argv[0]
        print
        print "Default for both savefilename and okfilename is \"~/.ceres\""
        print "If only one argument is given, both savefilename and okfilename will be set to that value"
        print
        sys.exit(0)

    if len(sys.argv)==2:
        savefilename=sys.argv[1]
        okfilename=sys.argv[1]
    else:
        if len(sys.argv)==3:
            savefilename=sys.argv[1]
            okfilename=sys.argv[2]
        else:
            savefilename=os.path.expanduser("~/.ceres")
            okfilename=savefilename

    a=sys.argv[0]
    b=os.path.dirname(a)
    c=os.path.abspath(b)
    d=c+"/../etc/ceresconfig.glade"
    config=Config(
        gtk.glade.XML(d,"Settings"),
        savefilename
        )
    config.setTempFileName(okfilename)
    config.readConfFile(okfilename);
    config.show()

<commit_msg>fix problem with preferences dialog.<commit_after>

import os, os.path, string, tempfile, traceback,gtk,gtk.glade,sys

#sys.path.insert(0,"/usr/lib/pygtk1/lib/python"+sys.version[:3]+"/site-packages/gtk-1.2")
#sys.path.insert(0,os.path.abspath(os.path.dirname(sys.argv[0]))+"/../lib/python"+sys.version[:3]+"/site-packages/gtk-1.2")
#sys.path.insert(0, PYGTK1PLACEMENT  +"/../lib/python"+sys.version[:3]+"/site-packages/gtk-1.2")

def lrstrip(dasstring):
    dasstring=string.lstrip(dasstring)
    dasstring=string.rstrip(dasstring)
    return dasstring

class Config:
    def __init__(self,xmltree,conffilename):
        self.xmltree=xmltree
        self.vars={}
        self.conffilename=conffilename
        self.readConfFile(conffilename)
        self.fileselectors={}

    def setTempFileName(self,tempfilename):
        self.tempfilename=tempfilename

    def readConfFile(self,filename):
        try:
            file=open(filename,"r")
        except:
            return
        while 1:
            line=""
            while line=="" or line=="\n" or line[0:1]=="#":
                line=file.readline()
                if line=="":
                    file.close()
                    return
                line=lrstrip(line)

            splitline=string.split(line,"=",1)
            splitline[0]=lrstrip(splitline[0])
            splitline[1]=lrstrip(splitline[1])

            if isinstance(self.xmltree.get_widget(splitline[0]),get.Entry):
                self.vars[splitline[0]]=splitline[1]
            else:
                self.vars[splitline[0]]=int(splitline[1])

    def writeConfFile(self,filename):
        try:
            file=open(filename,"w")
        except:
            print "Could not write file "+filename

        varnames=self.vars.keys()
        varnames.sort()
        for varname in varnames:
            if isinstance(self.xmltree.get_widget(varname), gtk.Entry):
                file.write(varname+"="+self.vars[varname]+"\n")
            else:
                file.write("%s=%d\n" % (varname,self.vars[varname]))

        try:
            file.close()
        except:
            print "Could not write file "+filename

    def getVar(self,confname):
        if self.vars.has_key(confname):
            return self.vars[confname]
        widg=self.xmltree.get_widget(confname)
        if widg!=None:
            if isinstance(widg,gtk.Entry):
                return widg.get_text()
            else:
                return widg.props.active

    def generalHandler(self,widg):
        name=widg.get_name()

        if type(widg)== gtk.Button:
            if name=="Ok":
                self.writeConfFile(self.tempfilename)
                gtk.main_quit()
                return
            if name=="Cancel":
                gtk.main_quit()
                return
            if name=="Save":
                self.writeConfFile(self.conffilename)
                self.writeConfFile(self.tempfilename)
                if self.getVar("close_window_when_pressing_save")==1:
                    gtk.main_quit()
                return

            if name[:7]=="browse+":
                widgnamestrip=string.split(name,"+",1)
                widgname=widgnamestrip[1]
                self.fileselectors[widgname]=gtk.glade.XML(
                    os.path.abspath(os.path.dirname(sys.argv[0]))+"/../etc/ceresconfig.glade",widgname
                    )
                self.fileselectors[widgname].signal_autoconnect(handle_dict)
                widg=self.fileselectors[widgname].get_widget(widgname)
                textentry=self.xmltree.get_widget(widgname[3:])
                widg.set_filename(textentry.get_text())
            else:
                widgnamestrip=string.split(name,"-",1)
                widgname="fs_"+widgnamestrip[1]
                widg=self.fileselectors[widgname].get_widget(widgname)
                if widgnamestrip[1]=="fsok":
                    filename=widg.get_filename()
                    if widgname!="fs_frequency_file":
                        filename=os.path.dirname(filename)+"/"
                    textentry=self.xmltree.get_widget(widgnamestrip[1])
                    textentry.set_text(filename)
                    self.vars[widgnamestrip[1]]=filename
                widg.destroy()
                del self.fileselectors[widgname]
        else:
            if isinstance(widg,gtk.Entry):
                widgtext=widg.get_text()
                if name[:-1]=="temporary_path_":
                    if widgtext[-1]!="/":
                        widgtext=widgtext+"/"
                        widg.set_text(widgtext)
                self.vars[name]=widgtext
            else:
                self.vars[name]=widg.props.active

    def destroy(self):
        gtk.main_quit()

    def show(self):
        self.xmltree.signal_autoconnect(handle_dict)
        varnames=self.vars.keys()
        for varname in varnames:
            widget=self.xmltree.get_widget(varname)
            if isinstance(widget,  gtk.GtkEntry):
                widget.set_text(self.vars[varname])
            else:
                widget.set_active(self.vars[varname])

        gtk.main()

def getVar(confname):
    try:
        return config.getVar(confname)
    except:
        traceback.print_exc(file=sys.stderr)

def generalHandler(obj):
    config.generalHandler(obj)

def destroy(obj):
    config.destroy()

def showConfig():
    try:
        tempfilename=tempfile.mktemp()
        config.writeConfFile(tempfilename)
        os.system(
            os.path.abspath(os.path.dirname(sys.argv[0]))
            + "/ceresconfig.py %s %s" % (config.conffilename,tempfilename)
            )
        config.readConfFile(config.conffilename)
        config.readConfFile(tempfilename)
        os.unlink(tempfilename)
    except:
        traceback.print_exc(file=sys.stderr)

handle_dict={
    "general":generalHandler,
    "on_Settings_destroy":destroy
    }

def ceresConfigStart(filename):
    global config
    a=sys.argv[0]
    b=os.path.dirname(a)
    c=os.path.abspath(b)
    d=c+"/../etc/ceresconfig.glade"
    config=Config(
        gtk.glade.XML(d,"Settings"),
        filename
        )

if __name__=="__main__":
    global config

    if len(sys.argv)>1 and (sys.argv[1]=="--help" or sys.argv[1]=="-h" or sys.argv[1]=="-help"):
        print
        print "Usage: %s [savefilename] [okfilename]" % sys.argv[0]
        print
        print "Default for both savefilename and okfilename is \"~/.ceres\""
        print "If only one argument is given, both savefilename and okfilename will be set to that value"
        print
        sys.exit(0)

    if len(sys.argv)==2:
        savefilename=sys.argv[1]
        okfilename=sys.argv[1]
    else:
        if len(sys.argv)==3:
            savefilename=sys.argv[1]
            okfilename=sys.argv[2]
        else:
            savefilename=os.path.expanduser("~/.ceres")
            okfilename=savefilename

    a=sys.argv[0]
    b=os.path.dirname(a)
    c=os.path.abspath(b)
    d=c+"/../etc/ceresconfig.glade"
    config=Config(
        gtk.glade.XML(d,"Settings"),
        savefilename
        )
    config.setTempFileName(okfilename)
    config.readConfFile(okfilename);
    config.show()

<|endoftext|>


<commit_before>from irc.bot import IRCDict, Channel, SingleServerIRCBot
import Queue
import threading
import logging
import socks
import socket

class TwitchIRCBot(SingleServerIRCBot):

    def __init__(self, master, worker_name, bot_name, server, access_token, port=6667, proxy=None):
        self.master = master
        self.worker_name = worker_name
        self.command = "!"+self.master.bot_name
        self.started = False
        self.is_connected = False

        if proxy is not None:
            logging.info('[%s] Proxy set: %s:%s', self.worker_name, proxy["address"], proxy["port"])
            socks.set_default_proxy(socks.HTTP, proxy["address"], proxy["port"])
            socket.socket = socks.socksocket

        SingleServerIRCBot.__init__(self, [(server, port, access_token)], bot_name, bot_name)

        # Channels set up
        self.channels = IRCDict()
        self.channel_join_queue = Queue.Queue()
        self.channel_list = []

        # Messages set up
        self.user_message_queue = Queue.Queue()

        logging.info('[%s] Chat worker bot initialized.', self.worker_name)

    def on_welcome(self, serv, event):
        self.is_connected = True
        if not self.started:
            logging.info('[%s] Connected to Twitch.tv IRC.', self.worker_name)
            # Start channel joining thread
            threading.Thread(target=self.channel_joiner, args=(serv,)).start()
            # Start message sending thread
            threading.Thread(target=self.message_sender, args=(serv,)).start()
            self.started = True
        # Welcome is a reconnect, rejoin all channels
        else:
            logging.info('[%s] Reconnected to Twitch.tv IRC.', self.worker_name)
            for channel in self.channel_list:
                self.channel_join_queue.put(channel)

    def on_disconnect(self, serv, event):
        self.is_connected = False
        logging.warning('[%s] Lost connection to Twitch.tv IRC. Attempting to reconnect...', self.worker_name)

    def on_pubmsg(self, serv, event):
        message = ''.join(event.arguments).strip()
        author = event.source.nick
        channel = event.target

        if message.startswith(self.command+" ") or message == self.command:
            self.master.process_message(self.worker_name, channel, author, message[len(self.command):].strip())

    # Thread for joining channels, capped at a limit of 50 joins per 15 seconds to follow twitch's restrictions
    # 50 joins / 15 seconds = Join up to 5 channels every 1.5 seconds
    def channel_joiner(self, serv):
        join_count = 0
        join_limit = 5
        while not self.channel_join_queue.empty() and join_count < join_limit:
            if self.is_connected:
                channel = self.channel_join_queue.get()
                self.channels[channel] = Channel()
                self.channel_list.append(channel)
                serv.join(channel)
                logging.info("[%s] Joining channel %s", self.worker_name, channel)
                join_count += 1
        threading.Timer(1.5, self.channel_joiner, args=(serv,)).start()

    # Thread for sending messages, capped at a limit of 20 messages per 30 seconds to follow twitch's restrictions
    # 20 messages / 30 seconds = Send 1 message every 1.5 seconds
    def message_sender(self, serv):

        if self.is_connected and self.master.message_center.has_message(self.worker_name):
            mdata = self.master.message_center.get_message(self.worker_name)
            message = mdata["message"]
            channel = mdata["channel"]

            if message != self.master.message_center.last_message: # Do not send the same message twice
                serv.privmsg(channel, message)
                self.master.message_center.last_message = message
                logging.info("[%s] %s %s: %s", self.worker_name, channel, self.master.bot_name, message)

        threading.Timer(1.5, self.message_sender, args=(serv,)).start()<commit_msg>Fix so that commands to ignore case<commit_after>from irc.bot import IRCDict, Channel, SingleServerIRCBot
import Queue
import threading
import logging
import socks
import socket

class TwitchIRCBot(SingleServerIRCBot):

    def __init__(self, master, worker_name, bot_name, server, access_token, port=6667, proxy=None):
        self.master = master
        self.worker_name = worker_name
        self.command = "!"+self.master.bot_name
        self.started = False
        self.is_connected = False

        if proxy is not None:
            logging.info('[%s] Proxy set: %s:%s', self.worker_name, proxy["address"], proxy["port"])
            socks.set_default_proxy(socks.HTTP, proxy["address"], proxy["port"])
            socket.socket = socks.socksocket

        SingleServerIRCBot.__init__(self, [(server, port, access_token)], bot_name, bot_name)

        # Channels set up
        self.channels = IRCDict()
        self.channel_join_queue = Queue.Queue()
        self.channel_list = []

        # Messages set up
        self.user_message_queue = Queue.Queue()

        logging.info('[%s] Chat worker bot initialized.', self.worker_name)

    def on_welcome(self, serv, event):
        self.is_connected = True
        if not self.started:
            logging.info('[%s] Connected to Twitch.tv IRC.', self.worker_name)
            # Start channel joining thread
            threading.Thread(target=self.channel_joiner, args=(serv,)).start()
            # Start message sending thread
            threading.Thread(target=self.message_sender, args=(serv,)).start()
            self.started = True
        # Welcome is a reconnect, rejoin all channels
        else:
            logging.info('[%s] Reconnected to Twitch.tv IRC.', self.worker_name)
            for channel in self.channel_list:
                self.channel_join_queue.put(channel)

    def on_disconnect(self, serv, event):
        self.is_connected = False
        logging.warning('[%s] Lost connection to Twitch.tv IRC. Attempting to reconnect...', self.worker_name)

    def on_pubmsg(self, serv, event):
        message = ''.join(event.arguments).strip()
        author = event.source.nick
        channel = event.target

        if message.lower().startswith(self.command.lower()+" ") or message.lower() == self.command.lower():
            self.master.process_message(self.worker_name, channel, author, message[len(self.command):].strip())

    # Thread for joining channels, capped at a limit of 50 joins per 15 seconds to follow twitch's restrictions
    # 50 joins / 15 seconds = Join up to 5 channels every 1.5 seconds
    def channel_joiner(self, serv):
        join_count = 0
        join_limit = 5
        while not self.channel_join_queue.empty() and join_count < join_limit:
            if self.is_connected:
                channel = self.channel_join_queue.get()
                self.channels[channel] = Channel()
                self.channel_list.append(channel)
                serv.join(channel)
                logging.info("[%s] Joining channel %s", self.worker_name, channel)
                join_count += 1
        threading.Timer(1.5, self.channel_joiner, args=(serv,)).start()

    # Thread for sending messages, capped at a limit of 20 messages per 30 seconds to follow twitch's restrictions
    # 20 messages / 30 seconds = Send 1 message every 1.5 seconds
    def message_sender(self, serv):

        if self.is_connected and self.master.message_center.has_message(self.worker_name):
            mdata = self.master.message_center.get_message(self.worker_name)
            message = mdata["message"]
            channel = mdata["channel"]

            if message != self.master.message_center.last_message: # Do not send the same message twice
                serv.privmsg(channel, message)
                self.master.message_center.last_message = message
                logging.info("[%s] %s %s: %s", self.worker_name, channel, self.master.bot_name, message)

        threading.Timer(1.5, self.message_sender, args=(serv,)).start()<|endoftext|>


<commit_before>import requests
from requests.exceptions import InvalidSchema, ConnectionError
from requests.exceptions import MissingSchema, Timeout, TooManyRedirects
from urllib.parse import urlparse, urlencode
import os
import tempfile
import magic
import logging
import math
import time
import sys

requests.packages.urllib3.disable_warnings()
logging.getLogger("requests").setLevel(logging.CRITICAL)
logging.getLogger("urllib3").setLevel(logging.CRITICAL)

class NetworkError(Exception):
    pass

class ConditionError(Exception):
    pass

class UrlError(ConditionError):
    pass

class StatusError(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

class Network(object):

    __allowed_schemata = set(['http', 'https'])

    @staticmethod
    def testLink(tr, journal, conf, session, acceptedTypes):
        log = logging.getLogger(__name__)

        log.debug("Fetching %s" % (tr.uri))
        log.debug("Data: %s" % (str(tr.data)))

        s = urlparse(tr.uri).scheme
        if s not in Network.__allowed_schemata:
            raise UrlError(s + " is not an allowed schema")

        accept = ",".join(acceptedTypes)
        log.debug("Accept header: %s" % (accept))

        tr.headers["User-Agent"] = conf.getProperty("agent")
        tr.headers["Accept"] = accept

        # if not allowed to send cookies or don't have any, then cookies are
        # None -> should be safe to use them; maybe filter which to use?

        log.debug("Timeout set to: %s" % (str(conf.getProperty("timeout"))))
        attempt = 0
        max_attempts = conf.getProperty("maxAttempts")
        while attempt < max_attempts:
            try:
                log.debug("Requesting")
                r = session.request(tr.method,
                                    tr.uri + Network.__gen_param(tr),
                                    headers=tr.headers,
                                    timeout=conf.getProperty("timeout"),
                                    cookies=tr.cookies,
                                    verify=conf.getProperty("verifyHttps"),
                                    stream=True)
            except TooManyRedirects as e:
                log.exception("Error while downloading: too many redirects", e)
                raise NetworkError("Too many redirects") from e
            except (ConnectionError, Timeout) as e:
                log.warn("Error while downloading: %s" % e)
                if (attempt + 1) < max_attempts:
                    log.info("Retrying!")
                    wait = math.pow(10, attempt)
                    time.sleep(wait)
                else:
                    raise NetworkError("%s attempts failed" %
                                       str(max_attempts)) from e
                attempt = attempt + 1
            else:
                return Network.__process_link(tr, r, journal, log)
        msg = "All %s attempts to get %s failed." % (str(max_attempts),
                                                     tr.uri)
        journal.foundDefect(tr.idno, "neterr", "Network error", msg, 0.9)
        raise NetworkError(msg)

    @staticmethod
    def __process_link(transaction, r, journal, log):
        transaction.status = r.status_code

        if r.status_code != requests.codes.ok:
            journal.foundDefect(transaction.srcId, "badlink", "Invalid link",
                                transaction.uri, 1.0)
            raise StatusError(r.status_code)

        if transaction.uri != r.url:
            log.debug("Redirection: %s -> %s" % (transaction.uri, r.url))
            transaction.changePrimaryUri(r.url)

        transaction.type = Network.__getCT(r, journal, transaction.idno)
        Network.__store_cookies(transaction, r.cookies, journal)

        return r

    @staticmethod
    def __getCT(r, journal, idno):
        if 'content-type' in r.headers:
            ct = r.headers['content-type']
        elif 'Content-Type' in r.headers:
            ct = r.headers['Content-Type']
        else:
            ct = ''

        if not ct.strip():
            journal.foundDefect(idno, "badtype", "Content-type empty", None,
                                0.5)

        if ';' in ct:  # text/html;charset=utf-8 -> text/html
            ct = ct.split(';')[0]

        return ct

    @staticmethod
    def getContent(transaction, conf, journal):
        log = logging.getLogger(__name__)
        try:
            with tempfile.NamedTemporaryFile(delete=False,
                                             prefix=conf.getProperty("tmpPrefix"),
                                             suffix=conf.getProperty("tmpSuffix"),
                                             dir=conf.getProperty("tmpDir", "/tmp/")) as tmp:
                Network.__download(transaction, conf, tmp, journal, log)
        except ConnectionError as e:
            log.debug("Connection error: %s" % (format(e)))
            journal.foundDefect(transaction.srcId, "badlink", "Invalid link",
                                transaction.uri, 1.0)
            raise NetworkError from e
        except Timeout as e:
            log.error("Timeout")
            journal.foundDefect(transaction.srcId, "timeout",
                                "Link timed out", transaction.uri, 0.9)
            raise NetworkError() from e
        else:
            match, mime = Network.__test_content_type(transaction.type,
                                                      transaction.file)
            if not match:
                journal.foundDefect(transaction.idno, "type-mishmash",
                                    "Declared content-type doesn't match detected one",
                                    "Declared " + transaction.type + ",detected " + mime,
                                    0.3)

    @staticmethod
    def __download(transaction, conf, tmp, journal, log):
            transaction.file = tmp.name
            log.info("Downloading %s" % transaction.uri)
            log.debug("Downloading chunks into %s" % tmp.name)
            limit = conf.getProperty("maxContentLength", sys.maxsize)
            # jde jen o to, ze pokud maxContentLength neni zadana, pak chceme,
            # aby ten limit byl > nez 10 000 000, pri None to hazi vyjimku
            chsize = min(limit, 10000000)
            for chunk in transaction.request.iter_content(chunk_size=chsize):
                if chunk:
                    tmp.write(chunk)
            log.debug("%s downloaded." % transaction.uri)

    @staticmethod
    def __gen_param(transaction):
        if (transaction.data is not None) and \
           (len(transaction.data) > 0) and \
           (transaction.method in set(['GET', 'HEAD'])):
            param = "?"+urlencode(transaction.data)
        else:
            param = ""
        return param

    @staticmethod
    def __store_cookies(transaction, cookies, journal):
        for cookie in cookies:
            name = cookie.name
            value = cookie.value
            path = cookie.path
            secure = cookie.secure
            httpOnly = cookie.has_nonstandard_attr('httpOnly')
            journal.gotCookie(transaction, name, value, secure, httpOnly, path)
        transaction.cookies = cookies

    @staticmethod
    def __test_content_type(ctype, fname):
        mime = magic.from_file(fname, mime=True)
        return (mime == ctype), mime
<commit_msg>Refactoring<commit_after>import requests
from requests.exceptions import InvalidSchema, ConnectionError
from requests.exceptions import MissingSchema, Timeout, TooManyRedirects
from urllib.parse import urlparse, urlencode
import os
import tempfile
import magic
import logging
import math
import time
import sys

requests.packages.urllib3.disable_warnings()
logging.getLogger("requests").setLevel(logging.CRITICAL)
logging.getLogger("urllib3").setLevel(logging.CRITICAL)

class NetworkError(Exception):
    pass

class ConditionError(Exception):
    pass

class UrlError(ConditionError):
    pass

class StatusError(Exception):

    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)

class Network(object):

    __allowed_schemata = set(['http', 'https'])

    @staticmethod
    def __check_schema(uri):
        s = urlparse(uri).scheme
        if s not in Network.__allowed_schemata:
            raise UrlError(s + " is not an allowed schema")

    @staticmethod
    def __backoff(attempt, max_attempts, log):
        if (attempt + 1) < max_attempts:
            log.info("Retrying!")
            wait = math.pow(10, attempt)
            time.sleep(wait)
        else:
            raise NetworkError("%s attempts failed" %
                               str(max_attempts)) from e

    @staticmethod
    def testLink(tr, journal, conf, session, acceptedTypes):
        log = logging.getLogger(__name__)

        log.debug("Fetching %s" % (tr.uri))
        log.debug("Data: %s" % (str(tr.data)))

        Network.__check_schema(tr.uri)

        accept = ",".join(acceptedTypes)
        log.debug("Accept header: %s" % (accept))

        tr.headers["User-Agent"] = conf.getProperty("agent")
        tr.headers["Accept"] = accept

        # if not allowed to send cookies or don't have any, then cookies are
        # None -> should be safe to use them; maybe filter which to use?

        log.debug("Timeout set to: %s" % (str(conf.getProperty("timeout"))))
        attempt = 0
        max_attempts = conf.getProperty("maxAttempts")
        while attempt < max_attempts:
            try:
                log.debug("Requesting")
                r = session.request(tr.method,
                                    tr.uri + Network.__gen_param(tr),
                                    headers=tr.headers,
                                    timeout=conf.getProperty("timeout"),
                                    cookies=tr.cookies,
                                    verify=conf.getProperty("verifyHttps"),
                                    stream=True)
            except TooManyRedirects as e:
                log.exception("Error while downloading: too many redirects", e)
                raise NetworkError("Too many redirects") from e
            except (ConnectionError, Timeout) as e:
                log.warn("Error while downloading: %s" % e)
                Network.__backoff(attempt, max_attempts, log)
                attempt = attempt + 1
            else:
                return Network.__process_link(tr, r, journal, log)
        msg = "All %s attempts to get %s failed." % (str(max_attempts),
                                                     tr.uri)
        journal.foundDefect(tr.idno, "neterr", "Network error", msg, 0.9)
        raise NetworkError(msg)

    @staticmethod
    def __process_link(transaction, r, journal, log):
        transaction.status = r.status_code

        if r.status_code != requests.codes.ok:
            journal.foundDefect(transaction.srcId, "badlink", "Invalid link",
                                transaction.uri, 1.0)
            raise StatusError(r.status_code)

        if transaction.uri != r.url:
            log.debug("Redirection: %s -> %s" % (transaction.uri, r.url))
            transaction.changePrimaryUri(r.url)

        transaction.type = Network.__getCT(r, journal, transaction.idno)
        Network.__store_cookies(transaction, r.cookies, journal)

        return r

    @staticmethod
    def __getCT(r, journal, idno):
        if 'content-type' in r.headers:
            ct = r.headers['content-type']
        elif 'Content-Type' in r.headers:
            ct = r.headers['Content-Type']
        else:
            ct = ''

        if not ct.strip():
            journal.foundDefect(idno, "badtype", "Content-type empty", None,
                                0.5)

        if ';' in ct:  # text/html;charset=utf-8 -> text/html
            ct = ct.split(';')[0]

        return ct

    @staticmethod
    def getContent(transaction, conf, journal):
        log = logging.getLogger(__name__)
        try:
            with tempfile.NamedTemporaryFile(delete=False,
                                             prefix=conf.getProperty("tmpPrefix"),
                                             suffix=conf.getProperty("tmpSuffix"),
                                             dir=conf.getProperty("tmpDir", "/tmp/")) as tmp:
                Network.__download(transaction, conf, tmp, journal, log)
        except ConnectionError as e:
            log.debug("Connection error: %s" % (format(e)))
            journal.foundDefect(transaction.srcId, "badlink", "Invalid link",
                                transaction.uri, 1.0)
            raise NetworkError from e
        except Timeout as e:
            log.error("Timeout")
            journal.foundDefect(transaction.srcId, "timeout",
                                "Link timed out", transaction.uri, 0.9)
            raise NetworkError() from e
        else:
            match, mime = Network.__test_content_type(transaction.type,
                                                      transaction.file)
            if not match:
                journal.foundDefect(transaction.idno, "type-mishmash",
                                    "Declared content-type doesn't match detected one",
                                    "Declared " + transaction.type + ",detected " + mime,
                                    0.3)

    @staticmethod
    def __download(transaction, conf, tmp, journal, log):
            transaction.file = tmp.name
            log.info("Downloading %s" % transaction.uri)
            log.debug("Downloading chunks into %s" % tmp.name)
            limit = conf.getProperty("maxContentLength", sys.maxsize)
            # jde jen o to, ze pokud maxContentLength neni zadana, pak chceme,
            # aby ten limit byl > nez 10 000 000, pri None to hazi vyjimku
            chsize = min(limit, 10000000)
            for chunk in transaction.request.iter_content(chunk_size=chsize):
                if chunk:
                    tmp.write(chunk)
            log.debug("%s downloaded." % transaction.uri)

    @staticmethod
    def __gen_param(transaction):
        if (transaction.data is not None) and \
           (len(transaction.data) > 0) and \
           (transaction.method in set(['GET', 'HEAD'])):
            param = "?"+urlencode(transaction.data)
        else:
            param = ""
        return param

    @staticmethod
    def __store_cookies(transaction, cookies, journal):
        for cookie in cookies:
            name = cookie.name
            value = cookie.value
            path = cookie.path
            secure = cookie.secure
            httpOnly = cookie.has_nonstandard_attr('httpOnly')
            journal.gotCookie(transaction, name, value, secure, httpOnly, path)
        transaction.cookies = cookies

    @staticmethod
    def __test_content_type(ctype, fname):
        mime = magic.from_file(fname, mime=True)
        return (mime == ctype), mime
<|endoftext|>


<commit_before>#
"""Manage access to the clients, including authentication"""

import logging
import sys

LOG = logging.getLogger(__name__)

class ClientCache<commit_msg>add cache method<commit_after>#
"""Manage access to the clients, including authentication"""

import logging
import sys

LOG = logging.getLogger(__name__)

class ClientCache(object):
    """ Descripter class for caching created client"""
    def __init__(self, factory):
        self.factory = factory
        self._handle = None

    def __get__(self, instance, owner):
        # Tell the ClientManager to login to keystone
        if self._handle is None:
            self._handle = self.factory(instance)
        return self._handle

<|endoftext|>


<commit_before>import time
import mmap
import os
import re

from fabric.api import run, env, task, local, sudo
from fabric.colors import green
from fabric.context_managers import prefix, cd, settings as fabconfig
from fabric.operations import put

from fabric_colors.environment import set_target_env

@set_target_env
def _uwsgi_chk_log():
    print("Checking if we have a log directory and file for our uwsgi process")
    cmd = "[ -d {0}/{1}_uwsgi ] && echo '1'".format(env.path_releases, env.project_name)
    result = run(cmd)
    return result

@set_target_env
def _uwsgi_mk_log():
    cmd = "mkdir -p {0}/{1}_uwsgi".format(env.path_releases, env.project_name)
    cmd2 = "touch {0}/{1}_uwsgi/{2}.log".format(env.path_releases, env.project_name, env.project_name)
    run(cmd)
    run(cmd2)

@set_target_env
def _uwsgi_status():
    print("Checking the status of our uswgi processes")
    result = run("if ps aux | grep -v grep | grep -i 'uwsgi' | grep '{0}' | grep 'uwsgi_{1}'; then echo 'OK'; else echo 'NO'; fi"\
            .format(env.project_name, env.target))
    if result.stdout != "NO":
        print("uwsgi is running!")
        return True
    else:
        print("uwsgi is dead!")
        return False

@set_target_env
def _uwsgi_start(newrelic=False):

    if newrelic:
        # If user provides a value here, we will restart uwsgi with the newrelic
        # command
        newrelic = True
    else:
        # Otherwise, we will check for its presence in fabsettings, failing
        # which we will finally default to False
        newrelic = env.newrelic

    if not _uwsgi_chk_log():
        _uwsgi_mk_log()

    if newrelic:
        print(green("Starting uwsgi with NEW_RELIC_CONFIG_FILE"))
        run("NEW_RELIC_CONFIG_FILE={0}/newrelic.ini \
                newrelic-admin run-program uwsgi --ini {1}/uwsgi_{2}.ini".format(env.path_release_current, env.project_path, env.target))
    else:
        print("Starting uwsgi")
        run("uwsgi --ini {0}/uwsgi_{1}.ini".format(env.project_path, env.target))

@set_target_env
def _uwsgi_restart(newrelic):
    if _uwsgi_status():
        print("Restarting uwsgi processes gracefully")
        run("kill -HUP `cat /tmp/{0}.pid-{1}`".format(env.project_name, env.webserver_port))
    else:
        print("uwsgi processes are already dead. executing a start.")
        _uwsgi_start(newrelic)

@set_target_env
def _uwsgi_restart_violent(newrelic):
    print("Restarting uwsgi processes violently")
    run("kill -INT `cat /tmp/{0}.pid-{1}`".format(env.project_name, env.webserver_port))
    time.sleep(1)
    _uwsgi_start(newrelic)

@task(default=True)
@set_target_env
def uwsgi(command, newrelic=False):
    """
    Usage: `fab -R all uwsgi:command`. command=start/stop/restart/restart_violent/status, newrelic(optional)=True/False.
    """
    with prefix(env.activate):
        with cd(env.path_release_current) and fabconfig(warn_only=True):
            if command == "start":
                print("Starting uwsgi processes")
                _uwsgi_start(newrelic)
            elif command == "stop":
                print("Stopping uwsgi processes")
                run("uwsgi --stop /tmp/{0}.pid-{1}".format(env.project_name, env.webserver_port))
            elif command == "restart":
                _uwsgi_restart(newrelic)
            elif command == "restart_violent":
                _uwsgi_restart_violent(newrelic)
            elif command == "status":
                _uwsgi_status()

@task
@set_target_env
def initializer():
    """
    Check if host has the current project's uwsgi initializer. If not, create it.
    """
    path_to_initializer = "/usr/lib/systemd/system/"
    initializer_name = env.project_name + "_" + env.target + ".service"
    path = path_to_initializer + initializer_name
    print("Checking if {0} is available".format(initializer_name))
    cmd = """if [ -e "{0}" ]; then echo 1; else echo ""; fi""".format(path)
    results = run(cmd)
    if results:
        print("{0} is already available in the host.".format(initializer_name))
        exit()

    template_file = os.path.dirname(__file__) + "/uwsgi.service.template"
    f = open(template_file)
    s = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    matches = re.findall(r'{[^{}]*}', s)
    unique_matches = (list(set(matches)))
    text = s[:]
    dic = {}
    for um in unique_matches:
        replacement_string = env.get(um[1:][:-1])
        #print("{0} to be replaced by {1}".format(um, replacement_string))
        dic[um] = replacement_string

    text = replace_all(text, dic)
    s.close()
    f.close()
    new_file = '{0}_{1}.service'.format(env.project_name, env.target)
    f_new = open(new_file, 'w+')
    f_new.write(text)
    f_new.close()
    put(new_file, path_to_initializer, use_sudo=True)
    sudo("systemctl enable {0}".format(new_file))
    local('rm {0}'.format(new_file))

def replace_all(text, dic):
    for i, j in dic.iteritems():
        text = text.replace(i, j)
    return text
<commit_msg>Set our PROJECT_HOME environment settings before uwsgi start command gets executed<commit_after>import time
import mmap
import os
import re

from fabric.api import run, env, task, local, sudo
from fabric.colors import green
from fabric.context_managers import prefix, cd, settings as fabconfig
from fabric.operations import put

from fabric_colors.environment import set_target_env

@set_target_env
def _uwsgi_chk_log():
    print("Checking if we have a log directory and file for our uwsgi process")
    cmd = "[ -d {0}/{1}_uwsgi ] && echo '1'".format(env.path_releases, env.project_name)
    result = run(cmd)
    return result

@set_target_env
def _uwsgi_mk_log():
    cmd = "mkdir -p {0}/{1}_uwsgi".format(env.path_releases, env.project_name)
    cmd2 = "touch {0}/{1}_uwsgi/{2}.log".format(env.path_releases, env.project_name, env.project_name)
    run(cmd)
    run(cmd2)

@set_target_env
def _uwsgi_status():
    print("Checking the status of our uswgi processes")
    result = run("if ps aux | grep -v grep | grep -i 'uwsgi' | grep '{0}' | grep 'uwsgi_{1}'; then echo 'OK'; else echo 'NO'; fi"\
            .format(env.project_name, env.target))
    if result.stdout != "NO":
        print("uwsgi is running!")
        return True
    else:
        print("uwsgi is dead!")
        return False

@set_target_env
def _uwsgi_start(newrelic=False):

    if newrelic:
        # If user provides a value here, we will restart uwsgi with the newrelic
        # command
        newrelic = True
    else:
        # Otherwise, we will check for its presence in fabsettings, failing
        # which we will finally default to False
        newrelic = env.newrelic

    if not _uwsgi_chk_log():
        _uwsgi_mk_log()

    if newrelic:
        print(green("Starting uwsgi with NEW_RELIC_CONFIG_FILE"))
        run("""
            export PROJECT_HOME={0};
            NEW_RELIC_CONFIG_FILE={1}/newrelic.ini \
                newrelic-admin run-program uwsgi --ini {2}/uwsgi_{3}.ini"""\
                .format(env.path_releases, env.path_release_current, env.project_path, env.target))
    else:
        print("Starting uwsgi")
        run("""
            export PROJECT_HOME={0};
            uwsgi --ini {1}/uwsgi_{2}.ini"""\
                    .format(env.path_releases, env.project_path, env.target))

@set_target_env
def _uwsgi_restart(newrelic):
    if _uwsgi_status():
        print("Restarting uwsgi processes gracefully")
        run("kill -HUP `cat /tmp/{0}.pid-{1}`".format(env.project_name, env.webserver_port))
    else:
        print("uwsgi processes are already dead. executing a start.")
        _uwsgi_start(newrelic)

@set_target_env
def _uwsgi_restart_violent(newrelic):
    print("Restarting uwsgi processes violently")
    run("kill -INT `cat /tmp/{0}.pid-{1}`".format(env.project_name, env.webserver_port))
    time.sleep(1)
    _uwsgi_start(newrelic)

@task(default=True)
@set_target_env
def uwsgi(command, newrelic=False):
    """
    Usage: `fab -R all uwsgi:command`. command=start/stop/restart/restart_violent/status, newrelic(optional)=True/False.
    """
    with prefix(env.activate):
        with cd(env.path_release_current) and fabconfig(warn_only=True):
            if command == "start":
                print("Starting uwsgi processes")
                _uwsgi_start(newrelic)
            elif command == "stop":
                print("Stopping uwsgi processes")
                run("uwsgi --stop /tmp/{0}.pid-{1}".format(env.project_name, env.webserver_port))
            elif command == "restart":
                _uwsgi_restart(newrelic)
            elif command == "restart_violent":
                _uwsgi_restart_violent(newrelic)
            elif command == "status":
                _uwsgi_status()

@task
@set_target_env
def initializer():
    """
    Check if host has the current project's uwsgi initializer. If not, create it.
    """
    path_to_initializer = "/usr/lib/systemd/system/"
    initializer_name = env.project_name + "_" + env.target + ".service"
    path = path_to_initializer + initializer_name
    print("Checking if {0} is available".format(initializer_name))
    cmd = """if [ -e "{0}" ]; then echo 1; else echo ""; fi""".format(path)
    results = run(cmd)
    if results:
        print("{0} is already available in the host.".format(initializer_name))
        exit()

    template_file = os.path.dirname(__file__) + "/uwsgi.service.template"
    f = open(template_file)
    s = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
    matches = re.findall(r'{[^{}]*}', s)
    unique_matches = (list(set(matches)))
    text = s[:]
    dic = {}
    for um in unique_matches:
        replacement_string = env.get(um[1:][:-1])
        #print("{0} to be replaced by {1}".format(um, replacement_string))
        dic[um] = replacement_string

    text = replace_all(text, dic)
    s.close()
    f.close()
    new_file = '{0}_{1}.service'.format(env.project_name, env.target)
    f_new = open(new_file, 'w+')
    f_new.write(text)
    f_new.close()
    put(new_file, path_to_initializer, use_sudo=True)
    sudo("systemctl enable {0}".format(new_file))
    local('rm {0}'.format(new_file))

def replace_all(text, dic):
    for i, j in dic.iteritems():
        text = text.replace(i, j)
    return text
<|endoftext|>