<commit_before>import logging
from optparse import make_option
import os
import pprint
import re
import shutil
import tempfile
import zipfile

from django.core.management.base import BaseCommand
from django.db import transaction

import path

import amo
import amo.utils
from files.models import nfd_str
from files.utils import extract_xpi, RDF
from versions.models import Version

log = logging.getLogger('z.migrations')
old_version = re.compile(r'\.1$')

def _log(msg):
    # Sigh. Vanishing log messages.  No really.
    print msg
    log.info(msg)

class Command(BaseCommand):
    help = ("One time script to fix broken jetpack XPIs. See bug 692524 "
            "and bug 693429")
    option_list = BaseCommand.option_list + (
        make_option('--dev', action='store_true',
                    dest='dev', help='Runs against dev versions.'),
    )

    @transaction.commit_manually
    def handle(self, *args, **options):
        success, fails, skipped = 0, 0, 0
        failed_ver_ids = []
        if options.get('dev'):
            _log('using dev IDs')
            ids = DEV_VERSION_IDS
        else:
            _log('using production IDs')
            ids = VERSION_IDS
        for version in Version.objects.filter(pk__in=ids):
            try:
                if not old_version.search(version.version):
                    _log('skipped: Unexpected version: %s [%s]'
                         % (version.version, version.pk))
                    transaction.rollback()
                    skipped += 1
                    continue
                # eg. 1.0.1 -> 1.0.2
                new_ver_str = old_version.sub('.2', version.version)
                _log('version: %s [%s] -> %s' % (version, version.pk,
                                                 new_ver_str))
                version.version = new_ver_str
                version.version_int = None  # recalculate on save
                version.save()
                for file_ in version.files.all():
                    tmp = tempfile.mkdtemp()
                    try:
                        xpi_dir = os.path.join(tmp, 'xpi_dir')
                        os.mkdir(xpi_dir)
                        extract_xpi(file_.file_path, xpi_dir)
                        xpi = os.path.join(tmp, 'xpi.zip')
                        new_xpi = fix_xpi(xpi, xpi_dir, new_ver_str)
                        replace_xpi(file_, new_xpi, version)
                    finally:
                        shutil.rmtree(tmp)
            except:
                failed_ver_ids.append(version.pk)
                transaction.rollback()
                log.exception(' ** rollback()')
                fails += 1
            else:
                transaction.commit()
                success += 1

        _log('These versions failed: %s' % pprint.pformat(failed_ver_ids))
        _log('SUCCESS: %s' % success)
        _log('FAILS: %s' % fails)
        _log('SKIPPED: %s' % skipped)
        _log('TOTAL: %s' % (success + fails + skipped))

def fix_xpi(xpi, xpi_dir, new_version):
    with open(os.path.join(xpi_dir, 'install.rdf')) as f:
        data = RDF(f.read())
    parent = (data.dom.documentElement
              .getElementsByTagName('Description')[0])
    for v in parent.getElementsByTagName('em:version'):
        parent.removeChild(v)
        v.unlink()
    elem = data.dom.createElement('em:version')
    elem.appendChild(data.dom.createTextNode(new_version))
    parent.appendChild(elem)
    # Replace install.rdf:
    with open(os.path.join(xpi_dir, 'install.rdf'), 'w') as f:
        f.write(str(data))

    zip = zipfile.ZipFile(xpi, 'w')
    for root, dirs, files in os.walk(xpi_dir):
        relroot = root.replace(xpi_dir, '')
        if relroot.startswith('/'):
            relroot = relroot[1:]
        for f in files:
            zip.write(os.path.join(root, f), os.path.join(relroot, f))
    zip.close()

    return path.path(xpi)

def replace_xpi(file_, new_xpi, version):
    old_filename = file_.filename
    old_filepath = file_.file_path
    # Now that we have a new version, make a new filename.
    file_.filename = file_.generate_filename(extension='.xpi')
    file_.hash = file_.generate_hash(new_xpi)
    file_.size = int(max(1, round(new_xpi.size / 1024, 0)))  # kB
    file_.save()
    destinations = [path.path(version.path_prefix)]
    if file_.status in amo.MIRROR_STATUSES:
        destinations.append(path.path(version.mirror_path_prefix))
    for dest in destinations:
        if not dest.exists():
            dest.makedirs()
        file_dest = dest / nfd_str(file_.filename)
        if os.path.exists(file_dest):
            backup_file(file_dest)
        new_xpi.copyfile(file_dest)
        _log('file: %s [%s] at %s -> %s at %s' % (old_filename,
                                                  file_.pk, old_filepath,
                                                  file_.filename,
                                                  file_dest))

def backup_file(file_path, base=None, tries=1):
    if not base:
        base = '%s.bak' % file_path
    if tries > 1:
        new_path = '%s.%s' % (base, tries)
    else:
        new_path = base
    if os.path.exists(new_path):
        return backup_file(file_path, base=base, tries=tries + 1)
    _log('backed up: %s -> %s' % (file_path, new_path))
    shutil.copy(file_path, new_path)

# The exact versions to fix:
VERSION_IDS = [1272303, 1271533, 1271532, 1271531, 1271530, 1271527, 1271526,
1271525, 1271523, 1271524, 1271520, 1271519, 1271518, 1271515, 1271514,
1271513, 1271511, 1271500, 1271499, 1271498, 1271497, 1271494, 1271493,
1271479, 1271436, 1271435, 1271434, 1271433, 1271432, 1271431, 1271430,
1271429, 1271428, 1271427, 1271426, 1271425, 1271424, 1271422, 1271421,
1271420, 1271419, 1271418, 1271417, 1271416, 1271414, 1271413, 1271412,
1271411, 1271410, 1271409, 1271408, 1271407, 1271406, 1271403, 1271404,
1271405, 1271402, 1271401, 1271400, 1271399, 1271398, 1271397, 1271396,
1271393, 1271394, 1271392, 1271391, 1271390, 1271389, 1271388, 1271387,
1271386, 1271385, 1271384, 1271382, 1271383, 1271381, 1271378, 1271379,
1271380, 1271377, 1271376, 1271375, 1271374, 1271373, 1271372, 1271371,
1271370, 1271369, 1271368, 1271367, 1271366, 1271365, 1271364, 1271363,
1271361, 1271360, 1271359, 1271356, 1271357, 1271358, 1271355, 1271354,
1271353, 1271352, 1271351, 1271350, 1271349, 1271348, 1271347, 1271345,
1271344, 1271346, 1271342, 1271341, 1271340, 1271337, 1271338, 1271339,
1271336, 1271335, 1271334, 1271331, 1271332, 1271333, 1271330, 1271329,
1271328, 1271327, 1271326, 1271325, 1271324, 1271323, 1271322, 1271320,
1271318, 1271319, 1271317, 1271316, 1271315, 1271314, 1271313, 1271312,
1271311, 1271309, 1271310, 1271308, 1271307, 1271306, 1271305, 1271304,
1271303, 1271302, 1271300, 1271301, 1271299, 1271297, 1271296, 1271295,
1271294, 1271293, 1271292, 1271291, 1271290, 1271289, 1271288, 1271287,
1271286, 1271285, 1271283]

DEV_VERSION_IDS = [1270479, 1270478, 1270481, 1270480, 1270482, 1270483,
1270484, 1270486, 1270487, 1270488, 1270489, 1270490, 1270491, 1270492,
1270493, 1270494, 1270496, 1270495, 1270498, 1270497, 1270500, 1270501,
1270502, 1270503, 1270504, 1270505, 1270506, 1270507, 1270508, 1270509,
1270510, 1270511, 1270512, 1270513, 1270514, 1270515, 1270516, 1270517,
1270519, 1270520, 1270521, 1270522, 1270523, 1270524, 1270525, 1270526,
1270527, 1270528, 1270529, 1270530, 1270531, 1270532, 1270533, 1270534,
1270536, 1270537, 1270539, 1270540, 1270541, 1270542, 1270543, 1270544,
1270545, 1270546, 1270547, 1270548, 1270549, 1270550, 1270551, 1270552,
1270553, 1270554, 1270555, 1270556, 1270557, 1270558, 1270559, 1270560,
1270562, 1270561, 1270564, 1270565, 1270566, 1270567, 1270568, 1270569,
1270572, 1270573, 1270574, 1270575, 1270576, 1270577, 1270578, 1270579,
1270580, 1270582, 1270583, 1270584, 1270585, 1270586, 1270587, 1270589,
1270590, 1270591, 1270592, 1270593, 1270594, 1270597, 1270599, 1270600,
1270601, 1270602, 1270603, 1270633, 1270634, 1270635, 1270636, 1270637,
1270639, 1270640]
<commit_msg>Using zip -r now to handle empty directories in the fixed xpi (bug 693429)<commit_after>import logging
from optparse import make_option
import os
import pprint
import re
import shutil
import subprocess
import tempfile

from django.core.management.base import BaseCommand
from django.db import transaction

import path

import amo
import amo.utils
from files.models import nfd_str
from files.utils import extract_xpi, RDF
from versions.models import Version

log = logging.getLogger('z.migrations')
old_version = re.compile(r'\.1$')

def _log(msg):
    # Sigh. Vanishing log messages.  No really.
    print msg
    log.info(msg)

class Command(BaseCommand):
    help = ("One time script to fix broken jetpack XPIs. See bug 692524 "
            "and bug 693429")
    option_list = BaseCommand.option_list + (
        make_option('--dev', action='store_true',
                    dest='dev', help='Runs against dev versions.'),
    )

    @transaction.commit_manually
    def handle(self, *args, **options):
        success, fails, skipped = 0, 0, 0
        failed_ver_ids = []
        if options.get('dev'):
            _log('using dev IDs')
            ids = DEV_VERSION_IDS
        else:
            _log('using production IDs')
            ids = VERSION_IDS
        for version in Version.objects.filter(pk__in=ids):
            try:
                if not old_version.search(version.version):
                    _log('skipped: Unexpected version: %s [%s]'
                         % (version.version, version.pk))
                    transaction.rollback()
                    skipped += 1
                    continue
                # eg. 1.0.1 -> 1.0.2
                new_ver_str = old_version.sub('.2', version.version)
                _log('version: %s [%s] -> %s' % (version, version.pk,
                                                 new_ver_str))
                version.version = new_ver_str
                version.version_int = None  # recalculate on save
                version.save()
                for file_ in version.files.all():
                    tmp = tempfile.mkdtemp()
                    try:
                        xpi_dir = os.path.join(tmp, 'xpi_dir')
                        os.mkdir(xpi_dir)
                        extract_xpi(file_.file_path, xpi_dir)
                        xpi = os.path.join(tmp, 'xpi.zip')
                        new_xpi = fix_xpi(xpi, xpi_dir, new_ver_str)
                        replace_xpi(file_, new_xpi, version)
                    finally:
                        shutil.rmtree(tmp)
            except:
                failed_ver_ids.append(version.pk)
                transaction.rollback()
                log.exception(' ** rollback()')
                fails += 1
            else:
                transaction.commit()
                success += 1

        _log('These versions failed: %s' % pprint.pformat(failed_ver_ids))
        _log('SUCCESS: %s' % success)
        _log('FAILS: %s' % fails)
        _log('SKIPPED: %s' % skipped)
        _log('TOTAL: %s' % (success + fails + skipped))

def fix_xpi(xpi, xpi_dir, new_version):
    with open(os.path.join(xpi_dir, 'install.rdf')) as f:
        data = RDF(f.read())
    parent = (data.dom.documentElement
              .getElementsByTagName('Description')[0])
    for v in parent.getElementsByTagName('em:version'):
        parent.removeChild(v)
        v.unlink()
    elem = data.dom.createElement('em:version')
    elem.appendChild(data.dom.createTextNode(new_version))
    parent.appendChild(elem)
    # Replace install.rdf:
    with open(os.path.join(xpi_dir, 'install.rdf'), 'w') as f:
        f.write(str(data))

    wd = os.getcwd()
    try:
        os.chdir(xpi_dir)
        subprocess.check_call(['zip', '-r', xpi] + os.listdir(os.getcwd()))
    finally:
        os.chdir(wd)
    return path.path(xpi)

def replace_xpi(file_, new_xpi, version):
    old_filename = file_.filename
    old_filepath = file_.file_path
    # Now that we have a new version, make a new filename.
    file_.filename = file_.generate_filename(extension='.xpi')
    file_.hash = file_.generate_hash(new_xpi)
    file_.size = int(max(1, round(new_xpi.size / 1024, 0)))  # kB
    file_.save()
    destinations = [path.path(version.path_prefix)]
    if file_.status in amo.MIRROR_STATUSES:
        destinations.append(path.path(version.mirror_path_prefix))
    for dest in destinations:
        if not dest.exists():
            dest.makedirs()
        file_dest = dest / nfd_str(file_.filename)
        if os.path.exists(file_dest):
            backup_file(file_dest)
        new_xpi.copyfile(file_dest)
        _log('file: %s [%s] at %s -> %s at %s' % (old_filename,
                                                  file_.pk, old_filepath,
                                                  file_.filename,
                                                  file_dest))

def backup_file(file_path, base=None, tries=1):
    if not base:
        base = '%s.bak' % file_path
    if tries > 1:
        new_path = '%s.%s' % (base, tries)
    else:
        new_path = base
    if os.path.exists(new_path):
        return backup_file(file_path, base=base, tries=tries + 1)
    _log('backed up: %s -> %s' % (file_path, new_path))
    shutil.copy(file_path, new_path)

# The exact versions to fix:
VERSION_IDS = [1272303, 1271533, 1271532, 1271531, 1271530, 1271527, 1271526,
1271525, 1271523, 1271524, 1271520, 1271519, 1271518, 1271515, 1271514,
1271513, 1271511, 1271500, 1271499, 1271498, 1271497, 1271494, 1271493,
1271479, 1271436, 1271435, 1271434, 1271433, 1271432, 1271431, 1271430,
1271429, 1271428, 1271427, 1271426, 1271425, 1271424, 1271422, 1271421,
1271420, 1271419, 1271418, 1271417, 1271416, 1271414, 1271413, 1271412,
1271411, 1271410, 1271409, 1271408, 1271407, 1271406, 1271403, 1271404,
1271405, 1271402, 1271401, 1271400, 1271399, 1271398, 1271397, 1271396,
1271393, 1271394, 1271392, 1271391, 1271390, 1271389, 1271388, 1271387,
1271386, 1271385, 1271384, 1271382, 1271383, 1271381, 1271378, 1271379,
1271380, 1271377, 1271376, 1271375, 1271374, 1271373, 1271372, 1271371,
1271370, 1271369, 1271368, 1271367, 1271366, 1271365, 1271364, 1271363,
1271361, 1271360, 1271359, 1271356, 1271357, 1271358, 1271355, 1271354,
1271353, 1271352, 1271351, 1271350, 1271349, 1271348, 1271347, 1271345,
1271344, 1271346, 1271342, 1271341, 1271340, 1271337, 1271338, 1271339,
1271336, 1271335, 1271334, 1271331, 1271332, 1271333, 1271330, 1271329,
1271328, 1271327, 1271326, 1271325, 1271324, 1271323, 1271322, 1271320,
1271318, 1271319, 1271317, 1271316, 1271315, 1271314, 1271313, 1271312,
1271311, 1271309, 1271310, 1271308, 1271307, 1271306, 1271305, 1271304,
1271303, 1271302, 1271300, 1271301, 1271299, 1271297, 1271296, 1271295,
1271294, 1271293, 1271292, 1271291, 1271290, 1271289, 1271288, 1271287,
1271286, 1271285, 1271283]

DEV_VERSION_IDS = [1270479, 1270478, 1270481, 1270480, 1270482, 1270483,
1270484, 1270486, 1270487, 1270488, 1270489, 1270490, 1270491, 1270492,
1270493, 1270494, 1270496, 1270495, 1270498, 1270497, 1270500, 1270501,
1270502, 1270503, 1270504, 1270505, 1270506, 1270507, 1270508, 1270509,
1270510, 1270511, 1270512, 1270513, 1270514, 1270515, 1270516, 1270517,
1270519, 1270520, 1270521, 1270522, 1270523, 1270524, 1270525, 1270526,
1270527, 1270528, 1270529, 1270530, 1270531, 1270532, 1270533, 1270534,
1270536, 1270537, 1270539, 1270540, 1270541, 1270542, 1270543, 1270544,
1270545, 1270546, 1270547, 1270548, 1270549, 1270550, 1270551, 1270552,
1270553, 1270554, 1270555, 1270556, 1270557, 1270558, 1270559, 1270560,
1270562, 1270561, 1270564, 1270565, 1270566, 1270567, 1270568, 1270569,
1270572, 1270573, 1270574, 1270575, 1270576, 1270577, 1270578, 1270579,
1270580, 1270582, 1270583, 1270584, 1270585, 1270586, 1270587, 1270589,
1270590, 1270591, 1270592, 1270593, 1270594, 1270597, 1270599, 1270600,
1270601, 1270602, 1270603, 1270633, 1270634, 1270635, 1270636, 1270637,
1270639, 1270640]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# VERSION: 0.2.1
# AUTHOR: Hunter Hammond
# DEPENDENCIES: pip install lxml

import argparse
from collections import OrderedDict
import itertools
import os
import re
import sys
import urllib2
import webbrowser

import lxml.html as lh

class CLIQuery:
    def __init__(self, url_args, search_flag, first_flag, open_flag, wolfram_flag):
        self.api_key, self.br_name = self.ReadConfig()
        self.search_flag = search_flag
        self.first_flag = first_flag
        self.open_flag = open_flag
        self.wolfram_flag = wolfram_flag
        self.url_args = self.ProcessArgs(url_args)
        self.html = self.GetHTML()

    def ReadConfig(self):
        script_dir = os.path.dirname(os.path.realpath(__file__))
        with open(script_dir + '/config.txt', 'r') as f:
            return f.readline().strip(), f.readline().strip()

    def ProcessArgs(self, url_args):
        clean_args = []
        if url_args:
            for arg in url_args:
                try:
                    if " " in arg:
                        for split_arg in arg.split(" "):
                            clean_args.append(split_arg)
                    else:
                        clean_args.append(arg)
                except IndexError:
                    sys.stderr.write('No search terms entered.\n')
                    sys.exit()
        else:
            sys.stderr.write('No search terms entered.\n')
            sys.exit()

        # Further processing of args before they are added to base_url
        new_url_args = []
        if not self.open_flag:
            try:
                for url_arg in clean_args:
                    if "+" in url_arg:
                        url_arg = arg.replace('+', '%2B')
                    new_url_args.append(url_arg)
                if len(new_url_args) > 1:
                    new_url_args = '+'.join(new_url_args)
                else:
                    new_url_args = new_url_args[0]
            except IndexError:
                sys.stderr.write('No search terms entered.\n')
                sys.exit()
        else:
            for url_arg in clean_args:
                if "." not in url_arg:
                    new_url_args.append(url_arg + ".com")
                else:
                    new_url_args.append(url_arg)
        return new_url_args

    def GetHTML(self):
        if not self.open_flag:
            if not self.wolfram_flag:
                return self.GetBingHTML(self.url_args)
            else:
                return self.GetWolframHTML(self.url_args)
        else:
            return ''

    def GetBingHTML(self, url_args):
        base_url = 'http://www.bing.com/search?q='
        url = base_url + url_args
        try:
            # Get HTML response
            request = urllib2.Request(url, headers={ 'User-Agent' : 'Mozilla/5.0' })
            return lh.parse(urllib2.urlopen(request))
        except urllib2.URLError:
            sys.stderr.write('Failed to retrieve webpage.\n')
            return ''

    def GetWolframHTML(self, url_args):
        base_url = 'http://api.wolframalpha.com/v2/query?input='
        url = base_url + url_args + '&appid=' + self.api_key
        try:
            # Get HTML response
            request = urllib2.Request(url, headers={ 'User-Agent' : 'Mozilla/5.0' })
            return lh.parse(urllib2.urlopen(request))
        except urllib2.URLError:
            sys.stderr.write('Failed to retrieve webpage.\n')
            return ''

    def BingSearch(self, html):
        # Parse Bing response and display link results
        unprocessed_links = html.xpath('//h2/a/@href')
        if not unprocessed_links:
            sys.stderr.write('Failed to retrieve webpage.\n')
            return False
        links = []
        link_descs = []
        for link in unprocessed_links:
            if not re.search('(ad|Ad|AD)(?=\W)', link): # Basic ad blocker
                if "http://" in link or "https://" in link:
                    links.append(link)
                    ld_xpath = "//h2/a[@href='" + str(link) + "']//text()"
                    link_desc = html.xpath(ld_xpath)
                    if type(link_desc) == list:
                        link_desc = ''.join(link_desc)
                    link_descs.append(link_desc)
                elif '/images/' in link and "www.bing.com" not in link:
                    # Fix image links that are not prepended with www.bing.com
                    links.append('http://www.bing.com' + link)
                    ld_xpath = "//h2/a[@href='" + str(link) + "']//text()"
                    link_desc = html.xpath(ld_xpath)
                    if type(link_desc) == list:
                        link_desc = ''.join(link_desc)
                    link_descs.append(link_desc)

        if links and link_descs:
            for i in xrange(len(links)):
                print_desc = (str(i) + ". " + link_descs[i]).encode('utf-8')
                print print_desc # Print link choices

            try:
                print ':',
                link_num = raw_input('')
                if link_num and int(link_num) >= 0 and int(link_num) < len(links):
                    self.OpenUrl(links[int(link_num)])
                    return True
            except (ValueError, IndexError):
                pass
        return False

    def WolframSearch(self, html):
        # Parse Wolfram|Alpha response for potential answers
        titles = list(OrderedDict.fromkeys(html.xpath("//pod[@title != '' and @title != 'Number line' and @title != 'Input' and @title != 'Visual representation' and @title != 'Image' and @title != 'Manipulatives illustration' and @title != 'Quotient and remainder']/@title")))
        entries = []
        if titles:
            for title in titles:
                entry_xpath = "//pod[@title='" + title + "']/subpod/plaintext/text()"
                entry = html.xpath(entry_xpath)
                if entry:
                    entries.append(entry[0])
            entries = list(OrderedDict.fromkeys(entries))
            output_list = []
            if len(entries) == 1 and entries[0] == '{}':
                return False
            for title, entry in itertools.izip(titles, entries):
                try:
                    if ' |' in entry:
                        entry = '\n\t' + entry.replace(' |', ':').replace('\n', '\n\t')
                    if title == 'Result':
                        output_list.append(entry).encode('utf-8')
                    else:
                        output_list.append(title + ': ' + entry).encode('utf-8')
                except (AttributeError, UnicodeEncodeError):
                    pass
            if not output_list:
                return False
            elif len(output_list) > 2:
                print '\n'.join(output_list[:2]).encode('utf-8')
                print 'See more? (y/n):',
                see_more = raw_input('')
                if see_more == 'y' or see_more == 'Y':
                    print '\n'.join(output_list[2:]).encode('utf-8')
            else:
                print '\n'.join(output_list).encode('utf-8')
            return True
        else:
            return False

    def BingCalculation(self, html):
        calc_result = html.xpath('//span[@id="rcTB"]/text()|//div[@class="b_focusTextMedium"]/text()|//p[@class="b_secondaryFocus df_p"]/text()|//div[@class="b_xlText b_secondaryText"]/text()|//input[@id="uc_rv"]/@value')
        define_result = html.xpath('//ol[@class="b_dList b_indent"]/li/div/text()')
        try:
            # Check if calculation result is present or age/date
            if calc_result:
                if len(calc_result) == 1:
                    print (calc_result[0]).encode('utf-8')
                else:
                    print ('\n'.join(calc_result)).encode('utf-8')
                return True
            # Check if calculation result is a definition
            elif define_result:
                if len(define_result) == 1:
                    print (define_result[0]).encode('utf-8')
                else:
                    print ('\n'.join(define_result)).encode('utf-8')
                return True
        except AttributeError:
            pass
        return False

    def OpenFirstLink(self, html):
        try:
            unprocessed_links = html.xpath('//h2/a/@href')
            for link in unprocessed_links:
                if not re.search('(ad|Ad|AD)(?=\W)', link): # Basic ad block
                    if "http://" in link or "https://" in link:
                        self.OpenUrl(link)
                        sys.exit()
                    elif '/images/' in link:
                        link = 'http://www.bing.com' + link
                        self.OpenUrl(link)
                        sys.exit()
        except IndexError:
            sys.stderr.write('Failed to retrieve webpage.\n')
        except AttributeError:
            sys.exit()

    def OpenUrl(self, url_args):
        try:
            br = webbrowser.get(self.br_name)
        except webbrowser.Error:
            sys.stderr.write('Could not locate runnable browser, make sure the browser path in config is correct.\n')
            sys.exit()
        if type(url_args) == list:
            for arg in url_args:
                br.open(arg)
        else:
            br.open(url_args)

    def Search(self):
        continue_search = False
        if self.open_flag:
            self.OpenUrl(self.url_args)
        elif self.search_flag:
            self.BingSearch(self.html)
        elif self.first_flag:
            self.OpenFirstLink(self.html)
        elif self.wolfram_flag:
            success = self.WolframSearch(self.html)
            if not success:
                continue_search = True
        else:
            continue_search = True

        if continue_search:
            # Defaults to BingCalculation, then WolframSearch, then BingSearch
            if self.wolfram_flag:
                bing_html = self.GetBingHTML(self.url_args)
                wolf_html = self.html
            else:
                bing_html = self.html
            if not self.BingCalculation(bing_html):
                if not self.wolfram_flag:
                    wolf_html = self.GetWolframHTML(self.url_args)
                if not self.WolframSearch(wolf_html):
                    self.BingSearch(bing_html)

if __name__ == "__main__":
    # parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--search", help="Get Bing search results",
    action="store_true")
    parser.add_argument("-f", "--first", help="Open first link result",
    action="store_true")
    parser.add_argument("-o", "--open", help="Open url directly",
    action="store_true")
    parser.add_argument("-w", "--wolfram", help="Get Wolfram|Alpha search results",
    action="store_true")
    parser.add_argument("URL_ARGS", nargs='*', help="Search keywords")
    args = parser.parse_args()
    search = bool(args.search)
    first = bool(args.first)
    openurl = bool(args.open)
    wolfram = bool(args.wolfram)
    query = CLIQuery(args.URL_ARGS, search, first, openurl, wolfram)
    query.Search()

<commit_msg>url to url_arg bug fix, added symbol dict<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# VERSION: 0.2.1
# AUTHOR: Hunter Hammond
# DEPENDENCIES: pip install lxml

import argparse
from collections import OrderedDict
import itertools
import os
import re
import sys
import urllib2
import webbrowser

import lxml.html as lh

class CLIQuery:
    def __init__(self, url_args, search_flag, first_flag, open_flag, wolfram_flag):
        self.api_key, self.br_name = self.ReadConfig()
        self.search_flag = search_flag
        self.first_flag = first_flag
        self.open_flag = open_flag
        self.wolfram_flag = wolfram_flag
        self.url_args = self.ProcessArgs(url_args)
        self.html = self.GetHTML()

    def ReadConfig(self):
        script_dir = os.path.dirname(os.path.realpath(__file__))
        with open(script_dir + '/config.txt', 'r') as f:
            return f.readline().strip(), f.readline().strip()

    def ProcessArgs(self, url_args):
        clean_args = []
        if url_args:
            for arg in url_args:
                try:
                    if " " in arg:
                        for split_arg in arg.split(" "):
                            clean_args.append(split_arg)
                    else:
                        clean_args.append(arg)
                except IndexError:
                    sys.stderr.write('No search terms entered.\n')
                    sys.exit()
        else:
            sys.stderr.write('No search terms entered.\n')
            sys.exit()

        # Further processing of args before they are added to base_url
        new_url_args = []
        if not self.open_flag:
            try:
                symbol_dict = { '@' : '%40',
                                '$' : '%24',
                                '%' : '%25',
                                '&' : '%26',
                                '+' : '%2B',
                                '=' : '%3D' }
                for url_arg in clean_args:
                    for sym in symbol_dict:
                        if sym in url_arg:
                            url_arg = url_arg.replace(sym, symbol_dict[sym])
                    new_url_args.append(url_arg)
                new_url_args = '+'.join(new_url_args)
            except IndexError:
                sys.stderr.write('No search terms entered.\n')
                sys.exit()
        else:
            for url_arg in clean_args:
                if "." not in url_arg:
                    new_url_args.append(url_arg + ".com")
                else:
                    new_url_args.append(url_arg)
        return new_url_args

    def GetHTML(self):
        if not self.open_flag:
            if not self.wolfram_flag:
                return self.GetBingHTML(self.url_args)
            else:
                sys.stderr.write(self.url_args)
                return self.GetWolframHTML(self.url_args)
        else:
            return ''

    def GetBingHTML(self, url_args):
        base_url = 'http://www.bing.com/search?q='
        url = base_url + url_args
        try:
            # Get HTML response
            request = urllib2.Request(url, headers={ 'User-Agent' : 'Mozilla/5.0' })
            return lh.parse(urllib2.urlopen(request))
        except urllib2.URLError:
            sys.stderr.write('Failed to retrieve webpage.\n')
            return ''

    def GetWolframHTML(self, url_args):
        base_url = 'http://api.wolframalpha.com/v2/query?input='
        url = base_url + url_args + '&appid=' + self.api_key
        try:
            # Get HTML response
            request = urllib2.Request(url, headers={ 'User-Agent' : 'Mozilla/5.0' })
            return lh.parse(urllib2.urlopen(request))
        except urllib2.URLError:
            sys.stderr.write('Failed to retrieve webpage.\n')
            return ''

    def BingSearch(self, html):
        # Parse Bing response and display link results
        try:
            unprocessed_links = html.xpath('//h2/a/@href')
        except AttributeError:
            sys.exit()
        if not unprocessed_links:
            sys.stderr.write('Failed to retrieve webpage.\n')
            return False
        links = []
        link_descs = []
        for link in unprocessed_links:
            if not re.search('(ad|Ad|AD)(?=\W)', link): # Basic ad blocker
                if "http://" in link or "https://" in link:
                    links.append(link)
                    ld_xpath = "//h2/a[@href='" + str(link) + "']//text()"
                    link_desc = html.xpath(ld_xpath)
                    if type(link_desc) == list:
                        link_desc = ''.join(link_desc)
                    link_descs.append(link_desc)
                elif '/images/' in link and "www.bing.com" not in link:
                    # Fix image links that are not prepended with www.bing.com
                    links.append('http://www.bing.com' + link)
                    ld_xpath = "//h2/a[@href='" + str(link) + "']//text()"
                    link_desc = html.xpath(ld_xpath)
                    if type(link_desc) == list:
                        link_desc = ''.join(link_desc)
                    link_descs.append(link_desc)

        if links and link_descs:
            for i in xrange(len(links)):
                print_desc = (str(i) + ". " + link_descs[i]).encode('utf-8')
                print print_desc # Print link choices

            try:
                print ':',
                link_num = raw_input('')
                if link_num and int(link_num) >= 0 and int(link_num) < len(links):
                    self.OpenUrl(links[int(link_num)])
                    return True
            except (ValueError, IndexError):
                pass
        return False

    def WolframSearch(self, html):
        # Parse Wolfram|Alpha response for potential answers
        try:
            titles = list(OrderedDict.fromkeys(html.xpath("//pod[@title != '' and @title != 'Number line' and @title != 'Input' and @title != 'Visual representation' and @title != 'Image' and @title != 'Manipulatives illustration' and @title != 'Quotient and remainder']/@title")))
        except AttributeError:
            sys.exit()
        entries = []
        if titles:
            for title in titles:
                entry_xpath = "//pod[@title='" + title + "']/subpod/plaintext/text()"
                entry = html.xpath(entry_xpath)
                if entry:
                    entries.append(entry[0])
            entries = list(OrderedDict.fromkeys(entries))
            output_list = []
            if len(entries) == 1 and entries[0] == '{}':
                return False
            for title, entry in itertools.izip(titles, entries):
                try:
                    if ' |' in entry:
                        entry = '\n\t' + entry.replace(' |', ':').replace('\n', '\n\t')
                    if title == 'Result':
                        output_list.append(entry).encode('utf-8')
                    else:
                        output_list.append(title + ': ' + entry).encode('utf-8')
                except (AttributeError, UnicodeEncodeError):
                    pass
            if not output_list:
                return False
            elif len(output_list) > 2:
                print '\n'.join(output_list[:2]).encode('utf-8')
                print 'See more? (y/n):',
                see_more = raw_input('')
                if see_more == 'y' or see_more == 'Y':
                    print '\n'.join(output_list[2:]).encode('utf-8')
            else:
                print '\n'.join(output_list).encode('utf-8')
            return True
        else:
            return False

    def BingCalculation(self, html):
        try:
            calc_result = html.xpath('//span[@id="rcTB"]/text()|//div[@class="b_focusTextMedium"]/text()|//p[@class="b_secondaryFocus df_p"]/text()|//div[@class="b_xlText b_secondaryText"]/text()|//input[@id="uc_rv"]/@value')
            define_result = html.xpath('//ol[@class="b_dList b_indent"]/li/div/text()')
        except AttributeError:
            sys.exit()
        try:
            # Check if calculation result is present or age/date
            if calc_result:
                if len(calc_result) == 1:
                    print (calc_result[0]).encode('utf-8')
                else:
                    print ('\n'.join(calc_result)).encode('utf-8')
                return True
            # Check if calculation result is a definition
            elif define_result:
                if len(define_result) == 1:
                    print (define_result[0]).encode('utf-8')
                else:
                    print ('\n'.join(define_result)).encode('utf-8')
                return True
        except AttributeError:
            pass
        return False

    def OpenFirstLink(self, html):
        try:
            unprocessed_links = html.xpath('//h2/a/@href')
            for link in unprocessed_links:
                if not re.search('(ad|Ad|AD)(?=\W)', link): # Basic ad block
                    if "http://" in link or "https://" in link:
                        self.OpenUrl(link)
                        sys.exit()
                    elif '/images/' in link:
                        link = 'http://www.bing.com' + link
                        self.OpenUrl(link)
                        sys.exit()
        except IndexError:
            sys.stderr.write('Failed to retrieve webpage.\n')
        except AttributeError:
            sys.exit()

    def OpenUrl(self, url_args):
        try:
            br = webbrowser.get(self.br_name)
        except webbrowser.Error:
            sys.stderr.write('Could not locate runnable browser, make sure the browser path in config is correct.\n')
            sys.exit()
        if type(url_args) == list:
            for arg in url_args:
                br.open(arg)
        else:
            br.open(url_args)

    def Search(self):
        continue_search = False
        if self.open_flag:
            self.OpenUrl(self.url_args)
        elif self.search_flag:
            self.BingSearch(self.html)
        elif self.first_flag:
            self.OpenFirstLink(self.html)
        elif self.wolfram_flag:
            success = self.WolframSearch(self.html)
            if not success:
                continue_search = True
        else:
            continue_search = True

        if continue_search:
            # Defaults to BingCalculation, then WolframSearch, then BingSearch
            if self.wolfram_flag:
                bing_html = self.GetBingHTML(self.url_args)
                wolf_html = self.html
            else:
                bing_html = self.html
            if not self.BingCalculation(bing_html):
                if not self.wolfram_flag:
                    wolf_html = self.GetWolframHTML(self.url_args)
                if not self.WolframSearch(wolf_html):
                    self.BingSearch(bing_html)

if __name__ == "__main__":
    # parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--search", help="Get Bing search results",
    action="store_true")
    parser.add_argument("-f", "--first", help="Open first link result",
    action="store_true")
    parser.add_argument("-o", "--open", help="Open url directly",
    action="store_true")
    parser.add_argument("-w", "--wolfram", help="Get Wolfram|Alpha search results",
    action="store_true")
    parser.add_argument("URL_ARGS", nargs='*', help="Search keywords")
    args = parser.parse_args()
    search = bool(args.search)
    first = bool(args.first)
    openurl = bool(args.open)
    wolfram = bool(args.wolfram)
    query = CLIQuery(args.URL_ARGS, search, first, openurl, wolfram)
    query.Search()

<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function

import numpy as np
import pandas as pd
from astropy.cosmology import Planck13 as cosmo
from astropy import units
import os

from cofm import c_DuttonMaccio as calc_c200
from smd_nfw import SurfaceMassDensity
import utils

try:
    from IPython.display import display
    notebook_display = True
except:
    notebook_display = False

class ClusterEnsemble(object):
    """Ensemble of galaxy clusters and their properties."""
    def __init__(self, redshifts):
        if type(redshifts) != np.ndarray:
            redshifts = np.array(redshifts)
        if redshifts.ndim != 1:
            raise ValueError("Input redshift array must have 1 dimension.")
        if np.sum(redshifts < 0.) > 0:
            raise ValueError("Redshifts cannot be negative.")
        self.describe = "Ensemble of galaxy clusters and their properties."
        self.number = redshifts.shape[0]
        self._z = redshifts
        self._rho_crit = cosmo.critical_density(self._z)
        self._massrich_norm = 2.7*10**13 * units.Msun
        self._massrich_slope = 1.4
        self._df = pd.DataFrame(self._z, columns=['z'])
        self._Dang_l = cosmo.angular_diameter_distance(self._z)
        self._m200 = None
        self._n200 = None
        self._r200 = None
        self._rs = None
        self._c200 = None
        self._deltac = None

    @property
    def n200(self):
        """Cluster richness values."""
        if self._n200 is None:
            raise AttributeError('n200 has not yet been initialized.')
        else:
            return self._n200

    @n200.setter
    def n200(self, richness):
        #Creates/updates values of cluster N200s & dependant variables.
        self._n200 = utils.check_units_and_type(richness, None,
                                                num = self.number)
        self._df['n200'] = pd.Series(self._n200, index = self._df.index)
        self._richness_to_mass()

    def _richness_to_mass(self):
        #Calculates M_200 for simple power-law scaling relation
        #(with default parameters from arXiv:1409.3571)
        self._m200 = self._massrich_norm * ((self._n200/20.) **
                                            self._massrich_slope)
        self._df['m200'] = pd.Series(self._m200, index = self._df.index)
        self._update_dependant_variables()

    @property
    def z(self):
        """Cluster redshifts."""
        return self._z

    @z.setter
    def z(self, redshifts):
        #Changes the values of the cluster z's and z-dependant variables.
        self._z = utils.check_units_and_type(redshifts, None,
                                             num = self.number)
        self._Dang_l = cosmo.angular_diameter_distance(self._z)
        self._df['z'] = pd.Series(self._z, index = self._df.index)
        self._rho_crit = cosmo.critical_density(self._z)
        if self._n200 is not None:
            self._update_dependant_variables()

    #note: user can access, but not modify, functions that are ONLY decorated
    # as properties; attempting to modify them will raise an AttributeError

    @property
    def Dang_l(self):
        return self._Dang_l

    @property
    def m200(self):
        if self._m200 is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._m200

    @property
    def dataframe(self):
        return self._df

    def _update_dependant_variables(self):
        self._calculate_r200()
        self._calculate_c200()
        self._calculate_rs()
        #what else depends on z or m or?

    # QUESTION: better design for following two functions?
    # could have a @property, @setter for each of norm and slope,
    # but what if user wants to modify both? would end up making
    # two calls to _richness_to_mass()... maybe thats no big deal?

    def massrich_parameters(self):
        """Print values of M200-N200 scaling relation parameters."""
        print("\nMass-Richness Power Law: M200 = norm * (N200 / 20) ^ slope")
        print("   norm:", self._massrich_norm)
        print("   slope:", self._massrich_slope)

    def update_massrichrelation(self, norm = None, slope = None):
        """Updates scaling relation parameters & dependant variables."""
        if norm is not None:
            if type(norm) == float:
                self._massrich_norm = norm * units.Msun
            else:
                raise TypeError("Input norm must be of type float, in units \
                of solar mass.")
        if slope is not None:
            if type(slope) == float:
                self._massrich_slope = slope
            else:
                raise TypeError("Input slope must be of type float.")
        if hasattr(self, 'n200'):
            self._richness_to_mass()

    def show(self, notebook = notebook_display):
        """Display table of cluster properties."""
        print("\nCluster Ensemble:")
        if notebook == True:
            display(self._df)
        elif notebook == False:
            print(self._df)
        self.massrich_parameters()

    @property
    def r200(self):
        if self._r200 is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._r200

    @property
    def c200(self):
        if self._c200 is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._c200

    @property
    def rs(self):
        if self._rs is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._rs

    @property
    def delta_c(self):
        if self._deltac is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._deltac

    def _calculate_r200(self):
        #calculate r200 from m200
        radius_200 = (3.*self._m200 / (800.*np.pi*self._rho_crit))**(1./3.)
        self._r200 = radius_200.to(units.Mpc)
        self._df['r200'] = pd.Series(self._r200, index = self._df.index)

    def _calculate_c200(self):
        #calculate c200 from m200 and z (using cofm.py)
        self._c200 = calc_c200(self._z,self._m200)
        self._df['c200'] = pd.Series(self._c200, index = self._df.index)
        self._calculate_deltac()

    def _calculate_rs(self):
        #cluster scale radius
        self._rs = self._r200 / self._c200
        self._df['rs'] = pd.Series(self._rs, index = self._df.index)

    def _calculate_deltac(self):
        #calculate concentration parameter from c200
        top = (200./3.)*self._c200**3.
        bottom = np.log(1.+self._c200)-(self._c200/(1.+self._c200))
        self._deltac = top/bottom
        self._df['delta_c'] = pd.Series(self._deltac, index = self._df.index)

    def calc_nfw(self, rbins, offsets = None, use_c = True):
        """Calculates Sigma and DeltaSigma NFW profiles of each cluster."""
        if offsets is None:
            self._sigoffset = np.zeros(self.number)*units.Mpc
        else:
            self._sigoffset = utils.check_units_and_type(offsets, units.Mpc,
                                                         num = self.number)

        self.rbins = utils.check_units_and_type(rbins, units.Mpc)

        if use_c:
            #--------
            #the old c way
            smdout = np.transpose(np.vstack(([self.rs],
                                            [self.delta_c],
                                            [self._rho_crit.to(units.Msun/
                                                            units.pc**3)],
                                            [self._sigoffset])))
            np.savetxt('smd_in1.dat', np.transpose(self.rbins), fmt='%15.8g')
            np.savetxt('smd_in2.dat', smdout, fmt='%15.8g')
            os.system('./smd_nfw')    #c program does the calculations
            sigma_nfw = np.loadtxt('sigma.dat')
            deltasigma_nfw = np.loadtxt('deltasigma.dat')
            os.system('rm -f smd_in1.dat')
            os.system('rm -f smd_in2.dat')
            os.system('rm -f sigma.dat')
            os.system('rm -f deltasigma.dat')
            #--------

            if offsets is None:
                self.sigma_nfw = sigma_nfw * units.Msun/(units.pc**2)
                self.deltasigma_nfw = deltasigma_nfw * units.Msun/(units.pc**2)
            else:
                self.sigma_offset = sigma_nfw * units.Msun/(units.pc**2)
                self.deltasigma_offset = deltasigma_nfw * units.Msun/(units.pc**2)

        else:
            #the python way
            smd = SurfaceMassDensity(self.rs,
                                     self.delta_c,
                                     self._rho_crit.to(units.Msun/units.pc**2/
                                                       units.Mpc),
                                     sig_offset = self._sigoffset,
                                     rbins = self.rbins)
            if offsets is None:
                self.sigma_nfw = smd.sigma_nfw()
                self.deltasigma_nfw = smd.deltasigma_nfw()
            else:
                #self.sigma_offset = smd.sigma_nfw()
                #self.deltasigma_offset = smd.deltasigma_nfw()
                #raise ValueError("Python does not yet calculate offset profiles. Use the c option.\n")
                self.sigma_nfw = smd.sigma_nfw(offsets=self._sigoffset)
                #TO DO: passing offsets is redundant with instantiation line 252
                # I think it should just be in instantiation, since thats where you
                # would be changing mass-related inputs in modeling...

<commit_msg>use same name for profiles, regardless of centered/offset<commit_after>from __future__ import absolute_import, division, print_function

import numpy as np
import pandas as pd
from astropy.cosmology import Planck13 as cosmo
from astropy import units
import os

from cofm import c_DuttonMaccio as calc_c200
from smd_nfw import SurfaceMassDensity
import utils

try:
    from IPython.display import display
    notebook_display = True
except:
    notebook_display = False

class ClusterEnsemble(object):
    """Ensemble of galaxy clusters and their properties."""
    def __init__(self, redshifts):
        if type(redshifts) != np.ndarray:
            redshifts = np.array(redshifts)
        if redshifts.ndim != 1:
            raise ValueError("Input redshift array must have 1 dimension.")
        if np.sum(redshifts < 0.) > 0:
            raise ValueError("Redshifts cannot be negative.")
        self.describe = "Ensemble of galaxy clusters and their properties."
        self.number = redshifts.shape[0]
        self._z = redshifts
        self._rho_crit = cosmo.critical_density(self._z)
        self._massrich_norm = 2.7*10**13 * units.Msun
        self._massrich_slope = 1.4
        self._df = pd.DataFrame(self._z, columns=['z'])
        self._Dang_l = cosmo.angular_diameter_distance(self._z)
        self._m200 = None
        self._n200 = None
        self._r200 = None
        self._rs = None
        self._c200 = None
        self._deltac = None

    @property
    def n200(self):
        """Cluster richness values."""
        if self._n200 is None:
            raise AttributeError('n200 has not yet been initialized.')
        else:
            return self._n200

    @n200.setter
    def n200(self, richness):
        #Creates/updates values of cluster N200s & dependant variables.
        self._n200 = utils.check_units_and_type(richness, None,
                                                num = self.number)
        self._df['n200'] = pd.Series(self._n200, index = self._df.index)
        self._richness_to_mass()

    def _richness_to_mass(self):
        #Calculates M_200 for simple power-law scaling relation
        #(with default parameters from arXiv:1409.3571)
        self._m200 = self._massrich_norm * ((self._n200/20.) **
                                            self._massrich_slope)
        self._df['m200'] = pd.Series(self._m200, index = self._df.index)
        self._update_dependant_variables()

    @property
    def z(self):
        """Cluster redshifts."""
        return self._z

    @z.setter
    def z(self, redshifts):
        #Changes the values of the cluster z's and z-dependant variables.
        self._z = utils.check_units_and_type(redshifts, None,
                                             num = self.number)
        self._Dang_l = cosmo.angular_diameter_distance(self._z)
        self._df['z'] = pd.Series(self._z, index = self._df.index)
        self._rho_crit = cosmo.critical_density(self._z)
        if self._n200 is not None:
            self._update_dependant_variables()

    #note: user can access, but not modify, functions that are ONLY decorated
    # as properties; attempting to modify them will raise an AttributeError

    @property
    def Dang_l(self):
        return self._Dang_l

    @property
    def m200(self):
        if self._m200 is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._m200

    @property
    def dataframe(self):
        return self._df

    def _update_dependant_variables(self):
        self._calculate_r200()
        self._calculate_c200()
        self._calculate_rs()
        #what else depends on z or m or?

    # QUESTION: better design for following two functions?
    # could have a @property, @setter for each of norm and slope,
    # but what if user wants to modify both? would end up making
    # two calls to _richness_to_mass()... maybe thats no big deal?

    def massrich_parameters(self):
        """Print values of M200-N200 scaling relation parameters."""
        print("\nMass-Richness Power Law: M200 = norm * (N200 / 20) ^ slope")
        print("   norm:", self._massrich_norm)
        print("   slope:", self._massrich_slope)

    def update_massrichrelation(self, norm = None, slope = None):
        """Updates scaling relation parameters & dependant variables."""
        if norm is not None:
            if type(norm) == float:
                self._massrich_norm = norm * units.Msun
            else:
                raise TypeError("Input norm must be of type float, in units \
                of solar mass.")
        if slope is not None:
            if type(slope) == float:
                self._massrich_slope = slope
            else:
                raise TypeError("Input slope must be of type float.")
        if hasattr(self, 'n200'):
            self._richness_to_mass()

    def show(self, notebook = notebook_display):
        """Display table of cluster properties."""
        print("\nCluster Ensemble:")
        if notebook == True:
            display(self._df)
        elif notebook == False:
            print(self._df)
        self.massrich_parameters()

    @property
    def r200(self):
        if self._r200 is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._r200

    @property
    def c200(self):
        if self._c200 is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._c200

    @property
    def rs(self):
        if self._rs is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._rs

    @property
    def delta_c(self):
        if self._deltac is None:
            raise AttributeError('Attribute has not yet been initialized.')
        else:
            return self._deltac

    def _calculate_r200(self):
        #calculate r200 from m200
        radius_200 = (3.*self._m200 / (800.*np.pi*self._rho_crit))**(1./3.)
        self._r200 = radius_200.to(units.Mpc)
        self._df['r200'] = pd.Series(self._r200, index = self._df.index)

    def _calculate_c200(self):
        #calculate c200 from m200 and z (using cofm.py)
        self._c200 = calc_c200(self._z,self._m200)
        self._df['c200'] = pd.Series(self._c200, index = self._df.index)
        self._calculate_deltac()

    def _calculate_rs(self):
        #cluster scale radius
        self._rs = self._r200 / self._c200
        self._df['rs'] = pd.Series(self._rs, index = self._df.index)

    def _calculate_deltac(self):
        #calculate concentration parameter from c200
        top = (200./3.)*self._c200**3.
        bottom = np.log(1.+self._c200)-(self._c200/(1.+self._c200))
        self._deltac = top/bottom
        self._df['delta_c'] = pd.Series(self._deltac, index = self._df.index)

    def calc_nfw(self, rbins, offsets = None, use_c = True):
        """Calculates Sigma and DeltaSigma NFW profiles of each cluster."""
        if offsets is None:
            self._sigoffset = np.zeros(self.number)*units.Mpc
        else:
            self._sigoffset = utils.check_units_and_type(offsets, units.Mpc,
                                                         num = self.number)

        self.rbins = utils.check_units_and_type(rbins, units.Mpc)

        if use_c:
            #--------
            #the old c way
            smdout = np.transpose(np.vstack(([self.rs],
                                            [self.delta_c],
                                            [self._rho_crit.to(units.Msun/
                                                            units.pc**3)],
                                            [self._sigoffset])))
            np.savetxt('smd_in1.dat', np.transpose(self.rbins), fmt='%15.8g')
            np.savetxt('smd_in2.dat', smdout, fmt='%15.8g')
            os.system('./smd_nfw')    #c program does the calculations
            sigma_nfw = np.loadtxt('sigma.dat')
            deltasigma_nfw = np.loadtxt('deltasigma.dat')
            os.system('rm -f smd_in1.dat')
            os.system('rm -f smd_in2.dat')
            os.system('rm -f sigma.dat')
            os.system('rm -f deltasigma.dat')
            #--------

            #if offsets is None:
            self.sigma_nfw = sigma_nfw * units.Msun/(units.pc**2)
            self.deltasigma_nfw = deltasigma_nfw * units.Msun/(units.pc**2)

        else:
            #the python way
            smd = SurfaceMassDensity(self.rs,
                                     self.delta_c,
                                     self._rho_crit.to(units.Msun/units.pc**2/
                                                       units.Mpc),
                                     sig_offset = self._sigoffset,
                                     rbins = self.rbins)

            self.sigma_nfw = smd.sigma_nfw()
            if offsets is None:
                self.deltasigma_nfw = smd.deltasigma_nfw() #not yet implemented for offsets
            #else:
                #raise ValueError("Python does not yet calculate offset profiles. Use the c option.\n")

<|endoftext|>


<commit_before>from ..errors import *

from .state import State
from .util import checkCompatibility
from .core import Context, Object, BoundObject, Type

#
# Variable
#
# A variable is a simple container for a value. The scope object may be used
# in conjunction with assignments and values for advanced functionality.

class Variable(BoundObject):
    type = None

    def __init__(self, name:str, type:Type = None, tokens = None):
        BoundObject.__init__(self, name, tokens)
        self.type = type

    def copy(self):
        return Variable(self.name, self.type, self.tokens)

    def verify(self):
        self.type.verify()

    def resolveType(self):
        return self.type

    def verifyAssignment(self, value:Object):
        if value is None:
            return

        if self.type is None:
            self.type = value.resolveType()

        self.verify()

        if not checkCompatibility(value.resolveType(), self.type):
            raise TypeError(message="Cannot assign").add(object=value).add(message="to").add(object=self)

    @property
    def local_context(self):
        return None

    def __repr__(self):
        return "{}:{}".format(self.name, self.type)
<commit_msg>Removed dead code<commit_after>from ..errors import *

from .state import State
from .util import checkCompatibility
from .core import Context, Object, BoundObject, Type

#
# Variable
#
# A variable is a simple container for a value. The scope object may be used
# in conjunction with assignments and values for advanced functionality.

class Variable(BoundObject):
    type = None

    def __init__(self, name:str, type:Type = None, tokens = None):
        BoundObject.__init__(self, name, tokens)
        self.type = type

    def verify(self):
        self.type.verify()

    def resolveType(self):
        return self.type

    def verifyAssignment(self, value:Object):
        if value is None:
            return

        if self.type is None:
            self.type = value.resolveType()

        self.verify()

        if not checkCompatibility(value.resolveType(), self.type):
            raise TypeError(message="Cannot assign").add(object=value).add(message="to").add(object=self)

    def __repr__(self):
        return "{}:{}".format(self.name, self.type)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import os
from flask import Blueprint, request, render_template, g
from config import config
from werkzeug import secure_filename

config['UPLOAD_FOLDER'] = 'static/img'
config['ALLOWED_EXTENSIONS'] = set(['png', 'jpg', 'jpeg', 'gif'])

cms_page = Blueprint('cms_page', __name__, template_folder='templates')

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in config['ALLOWED_EXTENSIONS']

# render catalogue with listed categories
@cms_page.route("/cms/")
def show_cms():
	return render_template("cms.html", editname="")

# render catalogue with product
@cms_page.route("/cms/<editname>")
def show_cms_editor(editname):
	if editname == "Remove Category":
		cats = read_categories()
		return render_template("cms.html", editname=editname, cats = cats)
	elif editname == "Add Product":
		cats = read_categories()
		return render_template("cms.html", editname=editname, cats = cats)
	elif editname == "Remove Product":
		prod = read_products()
		cats = read_categories()
		return render_template("cms.html", editname=editname, cats = cats, prod = prod)

	return render_template("cms.html", editname=editname)

@cms_page.route("/cms/Add Product", methods=['POST'])
def add_product():
	prodname = request.form['prodname']
	prodprice = request.form['prodprice']
	proddesc = request.form['proddesc']
	prodcat = request.form['prodcat']
	prodfile = request.files['prodfile']
	produrl = ""

	existing_products = read_products()
	cats = read_categories()
	if prodname in existing_products:
		return render_template("cms.html", editname = "Add Product", cats = cats, ins = "error")

	if prodfile and allowed_file(prodfile.filename) and not os.path.isfile(config['UPLOAD_FOLDER'] + "/" + prodfile.filename):
		# Make the filename safe, remove unsupported chars
		filename = secure_filename(prodfile.filename)
		prodfile.save(os.path.join(config['UPLOAD_FOLDER'], filename))
		produrl = config['UPLOAD_FOLDER'] + "/" + filename
	else:
		print "Image Upload Error"
	db = getattr(g, 'db', None)
	try:
		query = "INSERT INTO tbl_product (name, description, image_url, price, cat_id) VALUES (%s, %s, %s, %s, (SELECT id from tbl_category WHERE name=%s));"
		with db as cursor:
			data = (prodname, proddesc, produrl, prodprice, prodcat)
			cursor.execute(query, data)
			db.commit()
	except Exception as e:
		print e
	return render_template("cms.html", editname = "Add Product", cats = cats, ins = "success")

@cms_page.route("/cms/Remove Product", methods=['POST'])
def remove_product():
    prodname = request.form['removeprod']
    db = getattr(g, 'db', None)
    url = ""
    try:
        query = "SELECT image_url from tbl_product WHERE name = %s"
        with db as cursor:
            data = (prodname)
            cursor.execute(query, data)
            url = cursor.fetchone()[0]
    except Exception as e:
        print e
	if url:
	    os.remove(url)
    db = getattr(g, 'db', None)
    try:
        query = "DELETE FROM tbl_product WHERE name = %s"
        with db as cursor:
            data = (prodname)
            cursor.execute(query, data)
            db.commit()
    except Exception as e:
        print e
    p = read_products()
    return render_template("cms.html", editname = "Remove Product", prod = p, ins = "success")

@cms_page.route("/cms/Add Category", methods=['POST'])
def add_category():
	catname = request.form['catname']
	cats = read_categories()

	if catname in cats:
		return render_template("cms.html", editname = "Add Category", ins = "error")
	else:
		db = getattr(g, 'db', None)
		try:
			query = "insert into tbl_category (name) VALUES (%s);";
			with db as cursor:
				cursor.execute(query, catname)
				db.commit()
		except Exception as e:
			print e

	return render_template("cms.html", editname = "Add Category", ins = "success")

@cms_page.route("/cms/Remove Category", methods=['POST'])
def remove_category():
	catname = request.form['removecat']
	db = getattr(g, 'db', None)
	try:
		query = "DELETE FROM tbl_category WHERE name = %s;";
		with db as cursor:
			cursor.execute(query, catname)
			db.commit()
	except Exception as e:
		print e

	cats = read_categories()
	return render_template("cms.html", editname = "Remove Category", cats=cats, ins = "success")

# Read categories fom database and returm them as a list
def read_categories():
	db = getattr(g, 'db', None)
	cursor = db.cursor()
	cursor.execute("select name from tbl_category;")

	cats = []

	for x in cursor:
		cats.append(x[0])

	return cats

#Read products from catalogue where catname is correct
def read_products():
	db = getattr(g, 'db', None)
	cursor = db.cursor()
	query = "select name from tbl_product;"
	cursor.execute(query)
	prod = []

	for x in cursor.fetchall():
		prod.append(x[0])

	return prod<commit_msg>Cms page fix to allow db ops. Makes 4 spaces tabs.<commit_after># -*- coding: utf-8 -*-
import os
from flask import Blueprint, request, render_template, g
from config import config
from werkzeug import secure_filename

config['UPLOAD_FOLDER'] = 'static/img'
config['ALLOWED_EXTENSIONS'] = set(['png', 'jpg', 'jpeg', 'gif'])

cms_page = Blueprint('cms_page', __name__, template_folder='templates')

def allowed_file(filename):
	return '.' in filename and filename.rsplit('.', 1)[1] in config['ALLOWED_EXTENSIONS']

# render catalogue with listed categories
@cms_page.route("/cms/")
def show_cms():
	return render_template("cms.html", editname="")

# render catalogue with product
@cms_page.route("/cms/<editname>")
def show_cms_editor(editname):
	if editname == "Remove Category":
		cats = read_categories()
		return render_template("cms.html", editname=editname, cats = cats)
	elif editname == "Add Product":
		cats = read_categories()
		return render_template("cms.html", editname=editname, cats = cats)
	elif editname == "Remove Product":
		prod = read_products()
		cats = read_categories()
		return render_template("cms.html", editname=editname, cats = cats, prod = prod)

	return render_template("cms.html", editname=editname)

@cms_page.route("/cms/Add Product", methods=['POST'])
def add_product():
	prodname = request.form['prodname']
	prodprice = request.form['prodprice']
	proddesc = request.form['proddesc']
	prodcat = request.form['prodcat']
	prodfile = request.files['prodfile']
	produrl = ""

	existing_products = read_products()
	cats = read_categories()
	if prodname in existing_products:
		return render_template("cms.html", editname = "Add Product", cats = cats, ins = "error")

	if prodfile and allowed_file(prodfile.filename) and not os.path.isfile(config['UPLOAD_FOLDER'] + "/" + prodfile.filename):
		# Make the filename safe, remove unsupported chars
		filename = secure_filename(prodfile.filename)
		prodfile.save(os.path.join(config['UPLOAD_FOLDER'], filename))
		produrl = config['UPLOAD_FOLDER'] + "/" + filename
	else:
		print "Image Upload Error"
	db = getattr(g, 'db', None)
	try:
		query = "INSERT INTO tbl_product (name, description, image_url, price, cat_id) VALUES (%s, %s, %s, %s, (SELECT id from tbl_category WHERE name=%s));"
		with db as cursor:
			data = (prodname, proddesc, produrl, prodprice, prodcat)
			cursor.execute(query, data)
			db.commit()
	except Exception as e:
		print e
	return render_template("cms.html", editname = "Add Product", cats = cats, ins = "success")

@cms_page.route("/cms/Remove Product", methods=['POST'])
def remove_product():
	prodname = (request.form['removeprod'],)
	db = getattr(g, 'db', None)
	url = ""
	query = "SELECT image_url from tbl_product WHERE name = %s"
	with db as cursor:
		cursor.execute(query, prodname)
		url = cursor.fetchone()[0]
		if url:
			os.remove(url)

	query = "DELETE FROM tbl_product WHERE name = %s"
	with db as cursor:
		data = (prodname,)
		cursor.execute(query, data)
		db.commit()

	p = read_products()
	return render_template("cms.html", editname = "Remove Product", prod = p, ins = "success")

@cms_page.route("/cms/Add Category", methods=['POST'])
def add_category():
	catname = (request.form['catname'],)
	cats = read_categories()

	if catname[0] in cats:
		return render_template("cms.html", editname = "Add Category", ins = "error")
	else:
		db = getattr(g, 'db', None)
		query = "insert into tbl_category (name) VALUES (%s);";
		with db as cursor:
			cursor.execute(query, catname)
			db.commit()

	return render_template("cms.html", editname = "Add Category", ins = "success")

@cms_page.route("/cms/Remove Category", methods=['POST'])
def remove_category():
	catname = (request.form['removecat'],)
	db = getattr(g, 'db', None)
	try:
		query = "DELETE FROM tbl_category WHERE name = %s;";
		with db as cursor:
			cursor.execute(query, catname)
			db.commit()
	except Exception as e:
		print e

	cats = read_categories()
	return render_template("cms.html", editname = "Remove Category", cats=cats, ins = "success")

# Read categories fom database and returm them as a list
def read_categories():
	db = getattr(g, 'db', None)
	cursor = db.cursor()
	cursor.execute("select name from tbl_category;")

	cats = []

	for x in cursor:
		cats.append(x[0])

	return cats

#Read products from catalogue where catname is correct
def read_products():
	db = getattr(g, 'db', None)
	cursor = db.cursor()
	query = "select name from tbl_product;"
	cursor.execute(query)
	prod = []

	for x in cursor.fetchall():
		prod.append(x[0])

	return prod
<|endoftext|>


<commit_before><commit_msg>Fixed typo.<commit_after><|endoftext|>


<commit_before># Electron Cash - lightweight Bitcoin client
# Copyright (C) 2019 Axel Gembe <derago@gmail.com>
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import re
import subprocess
import sys
import threading
from enum import Enum

import stem.process
import stem.control

from .. import util
from ..util import PrintError
from ..utils import Event
from ..simple_config import SimpleConfig

if sys.platform in ('windows', 'win32'):
    _TOR_BINARY_NAME = os.path.join(
        os.path.dirname(__file__), '..', '..', 'tor.exe')
else:
    _TOR_BINARY_NAME = os.path.join(os.path.dirname(__file__), 'bin', 'tor')

_TOR_ENABLED_KEY = 'tor_enabled'
_TOR_ENABLED_DEFAULT = False

_TOR_SOCKS_PORT_KEY = 'tor_socks_port'
_TOR_SOCKS_PORT_DEFAULT = 0

class TorController(PrintError):
    class Status(Enum):
        STOPPING = 0
        STOPPED = 1
        STARTED = 2
        READY = 3

    _config: SimpleConfig = None
    _tor_process: subprocess.Popen = None
    _tor_read_thread: threading.Thread = None
    _tor_controller: stem.control.Controller = None

    status = Status.STOPPED
    status_changed = Event()

    active_socks_port: int = None
    active_control_port: int = None
    active_port_changed = Event()

    def __init__(self, config: SimpleConfig):
        if not config:
            raise AssertionError('TorController: config must be set')

        self._config = config

        socks_port = self._config.get(
            _TOR_SOCKS_PORT_KEY, _TOR_SOCKS_PORT_DEFAULT)
        if not socks_port or not self._check_port(int(socks_port)):
            # If no valid SOCKS port is set yet, we set the default
            self._config.set_key(_TOR_SOCKS_PORT_KEY, _TOR_SOCKS_PORT_DEFAULT)

    def __del__(self):
        self.status_changed.clear()
        self.active_port_changed.clear()

    # Opened Socks listener on 127.0.0.1:53544
    # Opened Control listener on 127.0.0.1:3300
    _listener_re = re.compile(r".*\[notice\] Opened ([^ ]*) listener on (.*)$")
    _endpoint_re = re.compile(r".*:(\d*)")

    def _tor_msg_handler(self, message: str):
        if util.is_verbose:
            self.print_msg(message)

        # Check if this is a "Opened listener" message and extract the information
        # into the active_socks_port and active_control_port variables
        listener_match = TorController._listener_re.match(message)
        if listener_match:
            listener_type = listener_match.group(1)
            listener_endpoint = listener_match.group(2)
            endpoint_match = TorController._endpoint_re.match(
                listener_endpoint)
            if endpoint_match:
                endpoint_port = int(endpoint_match.group(1))
                if listener_type == 'Socks':
                    self.active_socks_port = endpoint_port
                elif listener_type == 'Control':
                    self.active_control_port = endpoint_port
                    # The control port is the last port opened, so only notify after it
                    self.active_port_changed(self)

    def _read_tor_msg(self):
        while self._tor_process and not self._tor_process.poll():
            line = self._tor_process.stdout.readline().decode('utf-8', 'replace').strip()
            if not line:
                break
            self._tor_msg_handler(line)

    _orig_subprocess_popen = subprocess.Popen

    @staticmethod
    def _popen_monkey_patch(*args, **kwargs):
        if sys.platform in ('win32'):
            if hasattr(subprocess, 'CREATE_NO_WINDOW'):
                kwargs['creationflags'] = subprocess.CREATE_NO_WINDOW
            elif hasattr(subprocess, 'STARTUPINFO'):
                si = subprocess.STARTUPINFO()
                si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
                kwargs['startupinfo'] = si
        return TorController._orig_subprocess_popen(*args, **kwargs)

    def start(self):
        if self._tor_process:
            # Tor is already running
            return

        if not self.is_enabled():
            # Don't start Tor if not enabled
            return

        # When the socks port is set to zero, we let tor choose one
        socks_port = str(self.get_socks_port())
        if socks_port == '0':
            socks_port = 'auto'

        try:
            subprocess.Popen = TorController._popen_monkey_patch
            self._tor_process = stem.process.launch_tor_with_config(
                tor_cmd=_TOR_BINARY_NAME,
                completion_percent=0,  # We will monitor the bootstrap status
                init_msg_handler=self._tor_msg_handler,
                take_ownership=True,
                close_output=False,
                config={
                    'SocksPort': socks_port,
                    'ControlPort': 'auto',
                    'CookieAuthentication': '1',
                    'DataDirectory': os.path.join(self._config.path, 'tor'),
                    'Log': 'NOTICE stdout',
                },
            )
        except:
            self.print_exception("Failed to start Tor")
            self._tor_process = None
            return
        finally:
            subprocess.Popen = TorController._orig_subprocess_popen

        self._tor_read_thread = threading.Thread(
            target=self._read_tor_msg, name="Tor message reader")
        self._tor_read_thread.start()

        self.status = TorController.Status.STARTED
        self.status_changed(self)

        try:
            self._tor_controller = stem.control.Controller.from_port(
                port=self.active_control_port)
            self._tor_controller.authenticate()
            self._tor_controller.add_event_listener(
                self._handle_network_liveliness_event, stem.control.EventType.NETWORK_LIVENESS)
        except:
            self.print_exception("Failed to connect to Tor control port")
            self.stop()
            return

        self.print_error("started (Tor version {})".format(
            self._tor_controller.get_version()))

    def stop(self):
        if not self._tor_process:
            # Tor is not running
            return

        self.status = TorController.Status.STOPPING
        self.status_changed(self)

        self.active_socks_port = None
        self.active_control_port = None
        self.active_port_changed(self)

        if self._tor_controller:
            self._tor_controller.close()
            self._tor_controller = None

        if self._tor_process:
            self._tor_process.terminate()
            self._tor_process.wait()
            self._tor_process = None

        if self._tor_read_thread:
            self._tor_read_thread.join()
            self._tor_read_thread = None

        self.status = TorController.Status.STOPPED
        self.status_changed(self)
        self.print_error("stopped")

    def _handle_network_liveliness_event(self, event: stem.response.events.NetworkLivenessEvent):
        old_status = self.status
        self.status = TorController.Status.READY if event.status == 'UP' else TorController.Status.STARTED
        if old_status != self.status:
            self.status_changed(self)

    def _check_port(self, port: int) -> bool:
        if not isinstance(port, int):
            return False
        if port is None:
            return False
        if port != 0:  # Port 0 is automatic
            if port < 1024 or port > 65535:
                return False
        return True

    def set_enabled(self, enabled: bool):
        self._config.set_key(_TOR_ENABLED_KEY, enabled)
        if enabled:
            self.start()
        else:
            self.stop()

    def is_enabled(self) -> bool:
        return bool(self._config.get(_TOR_ENABLED_KEY, _TOR_ENABLED_DEFAULT))

    def set_socks_port(self, port: int):
        if not self._check_port(port):
            raise AssertionError('TorController: invalid port')

        self.stop()
        self._config.set_key(_TOR_SOCKS_PORT_KEY, port)
        self.start()

    def get_socks_port(self) -> int:
        socks_port = self._config.get(
            _TOR_SOCKS_PORT_KEY, _TOR_SOCKS_PORT_DEFAULT)
        if not self._check_port(int(socks_port)):
            raise AssertionError('TorController: invalid port')
        return int(socks_port)
<commit_msg>Tor: Monkey-patch Stem to never use a proxy for the control connection<commit_after># Electron Cash - lightweight Bitcoin client
# Copyright (C) 2019 Axel Gembe <derago@gmail.com>
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import re
import subprocess
import sys
import threading
import socket
from enum import Enum

import stem.socket
import stem.process
import stem.control

from .. import util
from ..util import PrintError
from ..utils import Event
from ..simple_config import SimpleConfig

# The network code monkey patches socket to use a SOCKS socket in Network.set_proxy.
# This causes the Tor control connection to go through Tor, which is not a good idea.
# To work around this, we monkey patch Stem to always create a standard socket for the
# control connection.
_orig_socket = socket.socket
def _make_socket_monkey_patch(self):
    try:
        control_socket = _orig_socket(socket.AF_INET, socket.SOCK_STREAM)
        control_socket.connect((self.address, self.port))
        return control_socket
    except socket.error as exc:
        raise stem.SocketError(exc)
stem.socket.ControlPort._make_socket = _make_socket_monkey_patch

if sys.platform in ('windows', 'win32'):
    _TOR_BINARY_NAME = os.path.join(
        os.path.dirname(__file__), '..', '..', 'tor.exe')
else:
    _TOR_BINARY_NAME = os.path.join(os.path.dirname(__file__), 'bin', 'tor')

_TOR_ENABLED_KEY = 'tor_enabled'
_TOR_ENABLED_DEFAULT = False

_TOR_SOCKS_PORT_KEY = 'tor_socks_port'
_TOR_SOCKS_PORT_DEFAULT = 0

class TorController(PrintError):
    class Status(Enum):
        STOPPING = 0
        STOPPED = 1
        STARTED = 2
        READY = 3

    _config: SimpleConfig = None
    _tor_process: subprocess.Popen = None
    _tor_read_thread: threading.Thread = None
    _tor_controller: stem.control.Controller = None

    status = Status.STOPPED
    status_changed = Event()

    active_socks_port: int = None
    active_control_port: int = None
    active_port_changed = Event()

    def __init__(self, config: SimpleConfig):
        if not config:
            raise AssertionError('TorController: config must be set')

        self._config = config

        socks_port = self._config.get(
            _TOR_SOCKS_PORT_KEY, _TOR_SOCKS_PORT_DEFAULT)
        if not socks_port or not self._check_port(int(socks_port)):
            # If no valid SOCKS port is set yet, we set the default
            self._config.set_key(_TOR_SOCKS_PORT_KEY, _TOR_SOCKS_PORT_DEFAULT)

    def __del__(self):
        self.status_changed.clear()
        self.active_port_changed.clear()

    # Opened Socks listener on 127.0.0.1:53544
    # Opened Control listener on 127.0.0.1:3300
    _listener_re = re.compile(r".*\[notice\] Opened ([^ ]*) listener on (.*)$")
    _endpoint_re = re.compile(r".*:(\d*)")

    def _tor_msg_handler(self, message: str):
        if util.is_verbose:
            self.print_msg(message)

        # Check if this is a "Opened listener" message and extract the information
        # into the active_socks_port and active_control_port variables
        listener_match = TorController._listener_re.match(message)
        if listener_match:
            listener_type = listener_match.group(1)
            listener_endpoint = listener_match.group(2)
            endpoint_match = TorController._endpoint_re.match(
                listener_endpoint)
            if endpoint_match:
                endpoint_port = int(endpoint_match.group(1))
                if listener_type == 'Socks':
                    self.active_socks_port = endpoint_port
                elif listener_type == 'Control':
                    self.active_control_port = endpoint_port
                    # The control port is the last port opened, so only notify after it
                    self.active_port_changed(self)

    def _read_tor_msg(self):
        while self._tor_process and not self._tor_process.poll():
            line = self._tor_process.stdout.readline().decode('utf-8', 'replace').strip()
            if not line:
                break
            self._tor_msg_handler(line)

    _orig_subprocess_popen = subprocess.Popen

    @staticmethod
    def _popen_monkey_patch(*args, **kwargs):
        if sys.platform in ('win32'):
            if hasattr(subprocess, 'CREATE_NO_WINDOW'):
                kwargs['creationflags'] = subprocess.CREATE_NO_WINDOW
            elif hasattr(subprocess, 'STARTUPINFO'):
                si = subprocess.STARTUPINFO()
                si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
                kwargs['startupinfo'] = si
        return TorController._orig_subprocess_popen(*args, **kwargs)

    def start(self):
        if self._tor_process:
            # Tor is already running
            return

        if not self.is_enabled():
            # Don't start Tor if not enabled
            return

        # When the socks port is set to zero, we let tor choose one
        socks_port = str(self.get_socks_port())
        if socks_port == '0':
            socks_port = 'auto'

        try:
            subprocess.Popen = TorController._popen_monkey_patch
            self._tor_process = stem.process.launch_tor_with_config(
                tor_cmd=_TOR_BINARY_NAME,
                completion_percent=0,  # We will monitor the bootstrap status
                init_msg_handler=self._tor_msg_handler,
                take_ownership=True,
                close_output=False,
                config={
                    'SocksPort': socks_port,
                    'ControlPort': 'auto',
                    'CookieAuthentication': '1',
                    'DataDirectory': os.path.join(self._config.path, 'tor'),
                    'Log': 'NOTICE stdout',
                },
            )
        except:
            self.print_exception("Failed to start Tor")
            self._tor_process = None
            return
        finally:
            subprocess.Popen = TorController._orig_subprocess_popen

        self._tor_read_thread = threading.Thread(
            target=self._read_tor_msg, name="Tor message reader")
        self._tor_read_thread.start()

        self.status = TorController.Status.STARTED
        self.status_changed(self)

        try:
            self._tor_controller = stem.control.Controller.from_port(
                port=self.active_control_port)
            self._tor_controller.authenticate()
            self._tor_controller.add_event_listener(
                self._handle_network_liveliness_event, stem.control.EventType.NETWORK_LIVENESS)
        except:
            self.print_exception("Failed to connect to Tor control port")
            self.stop()
            return

        self.print_error("started (Tor version {})".format(
            self._tor_controller.get_version()))

    def stop(self):
        if not self._tor_process:
            # Tor is not running
            return

        self.status = TorController.Status.STOPPING
        self.status_changed(self)

        self.active_socks_port = None
        self.active_control_port = None
        self.active_port_changed(self)

        if self._tor_controller:
            self._tor_controller.close()
            self._tor_controller = None

        if self._tor_process:
            self._tor_process.terminate()
            self._tor_process.wait()
            self._tor_process = None

        if self._tor_read_thread:
            self._tor_read_thread.join()
            self._tor_read_thread = None

        self.status = TorController.Status.STOPPED
        self.status_changed(self)
        self.print_error("stopped")

    def _handle_network_liveliness_event(self, event: stem.response.events.NetworkLivenessEvent):
        old_status = self.status
        self.status = TorController.Status.READY if event.status == 'UP' else TorController.Status.STARTED
        if old_status != self.status:
            self.status_changed(self)

    def _check_port(self, port: int) -> bool:
        if not isinstance(port, int):
            return False
        if port is None:
            return False
        if port != 0:  # Port 0 is automatic
            if port < 1024 or port > 65535:
                return False
        return True

    def set_enabled(self, enabled: bool):
        self._config.set_key(_TOR_ENABLED_KEY, enabled)
        if enabled:
            self.start()
        else:
            self.stop()

    def is_enabled(self) -> bool:
        return bool(self._config.get(_TOR_ENABLED_KEY, _TOR_ENABLED_DEFAULT))

    def set_socks_port(self, port: int):
        if not self._check_port(port):
            raise AssertionError('TorController: invalid port')

        self.stop()
        self._config.set_key(_TOR_SOCKS_PORT_KEY, port)
        self.start()

    def get_socks_port(self) -> int:
        socks_port = self._config.get(
            _TOR_SOCKS_PORT_KEY, _TOR_SOCKS_PORT_DEFAULT)
        if not self._check_port(int(socks_port)):
            raise AssertionError('TorController: invalid port')
        return int(socks_port)
<|endoftext|>