<commit_before>#-------------------------------------------------------------------------------------------------------------
## Import Statements
#-------------------------------------------------------------------------------------------------------------
from __future__ import print_function
import sys
import os
import subprocess as sp
import glob
import traceback
from .utils import conda_root
from .utils import get_species
from .utils import get_ggd_channels
from .utils import get_channel_data
from .utils import get_channeldata_url
from .utils import bypass_satsolver_on_install
from .utils import active_conda_env
from .show_env import activate_enviroment_variables
from .search import load_json, load_json_from_url, search_packages
from .uninstall import remove_from_condaroot, check_for_installation
from .utils import get_required_conda_version

SPECIES_LIST = get_species()
#-------------------------------------------------------------------------------------------------------------
## Argument Parser
#-------------------------------------------------------------------------------------------------------------
def add_install(p):
    c = p.add_parser('install', help="install a data recipe from ggd")
    c.add_argument("-c", "--channel", default="genomics", choices=get_ggd_channels(),
                     help="The ggd channel the desired recipe is stored in. (Default = genomics)")
    c.add_argument("-v", "--version", default="-1", help="A specific ggd package version to install. If the -v flag is not used the latest version will be installed.")
    c.add_argument("-d", "--debug", action="store_true", help="(Optional) When the -d flag is set debug output will be printed to stdout.")
    c.add_argument("name", help="the name of the recipe to install")
    c.set_defaults(func=install)

#-------------------------------------------------------------------------------------------------------------
## Functions/Methods
#-------------------------------------------------------------------------------------------------------------

def check_ggd_recipe(ggd_recipe,ggd_channel):
    """Method used to check if the desired package is in the ggd repo using the repo metadata file

    check_ggd_recipe
    ================
    Method to check if the ggd recipe exists. Uses search_packages from search.py to
     search the ggd-channel json file. If the recipe exists within the json file,
     the installation proceeds. If not, the instalation stops
    """

    CHANNEL_DATA_URL = get_channeldata_url(ggd_channel)
    jdict = load_json_from_url(CHANNEL_DATA_URL)
    package_list = [x[0] for x in search_packages(jdict, ggd_recipe)]
    if ggd_recipe in package_list:
        print("\n\t-> %s exists in the ggd-%s channel" %(ggd_recipe,ggd_channel))
        return(jdict)
    else:
        print("\n\t-> '%s' was not found in ggd-%s" %(ggd_recipe, ggd_channel))
        print("\t-> You can search for recipes using the ggd search tool: \n\t\t'ggd search -t %s'\n" %ggd_recipe)
        return(None)

def check_if_installed(ggd_recipe,ggd_jdict,ggd_version):
    """Method to check if the recipe has already been installed and is in the conda ggd storage path.

    check_if_installed
    ==================
    This method is used to check if the ggd package has been installed and is located in the ggd storage path.
     If it is already installed the installation stops. If it is not detected then installation continues.
    """

    species = ggd_jdict["packages"][ggd_recipe]["identifiers"]["species"]
    build = ggd_jdict["packages"][ggd_recipe]["identifiers"]["genome-build"]
    version = ggd_jdict["packages"][ggd_recipe]["version"]

    CONDA_ROOT = conda_root()
    path = os.path.join(CONDA_ROOT,"share","ggd",species,build,ggd_recipe,version)
    recipe_exists = glob.glob(path)
    if recipe_exists:
        # if the ggd_version designated to be installed does not match the installed version, install the designated version
        if ggd_version != version and ggd_version != "-1":
            return(False)
        else:
            print("\n\t-> '%s' is already installed." %ggd_recipe)
            print("\t-> You can find %s here: %s" %(ggd_recipe,path))
            sys.exit()
    else:
        print("\n\t-> %s version %s is not installed on your system" %(ggd_recipe, version))
        return(False)

def check_conda_installation(ggd_recipe,ggd_version):
    """Method used to check if the recipe has been installed using conda.

    check_conda_installation
    ========================
    This method is used to check if the ggd data package has been installed by the conda system,
     without being installed by the ggd system. If it has, the recipe needs to be uninstalled and
     reinstalled. If not, the system continues to install the package using ggd.
    """

    conda_package_list = sp.check_output(["conda", "list"]).decode('utf8')
    recipe_find = conda_package_list.find(ggd_recipe)
    if recipe_find == -1:
        print("\n\t-> %s has not been installed by conda" %ggd_recipe)
        return(False)
    else:
        if ggd_version != "-1": ## Check if ggd version was designated
            installed_version = conda_package_list[recipe_find:recipe_find+100].split("\n")[0].replace(" ","")[len(ggd_recipe)]
            if installed_version != ggd_version:
                print("\n\t-> %s version %s has not been installed by conda" %(ggd_recipe,str(ggd_version)))
                return(False)
            else:
                print("\n\t-> %s version %s has been installed by conda on your system and must be uninstalled to proceed." %(ggd_recipe,str(ggd_version)))
                print("\t-> To reinstall run:\n\t\t $ ggd uninstall %s \n\t\t $ ggd install %s" %(ggd_recipe,ggd_recipe))
                sys.exit()
        else: ## If version is not specified check if exact package in conda list
            start_index = conda_package_list.find(ggd_recipe)
            end_index = start_index + len(ggd_recipe)
            ## Check if it really is the package, or something similar to the package
            if conda_package_list[int(start_index) - 1] == "\n" and conda_package_list[int(end_index) + 1] == " ":
                print("\n\t-> %s has been installed by conda on your system and must be uninstalled to proceed." %ggd_recipe)
                print("\t-> To reinstall run:\n\t\t $ ggd uninstall %s \n\t\t $ ggd install %s" %(ggd_recipe,ggd_recipe))
                sys.exit()
            else: ## IF not exactly in conda list
                print("\n\t-> %s has not been installed by conda" %ggd_recipe)
                return(False)

def check_S3_bucket(ggd_recipe, ggd_jdict):
    """Method to check if the recipe is stored on the ggd S3 bucket. If so it installs from S3

    check_S3_bucket
    ==============
    This method is used to check if the recipe has been cached on aws S3 bucket. It returns true if the
     the recipe is cached, and false if it is not. If it is cached the cached version will be installed.
    """

    if "tags" in ggd_jdict["packages"][ggd_recipe]:
        if "cached" in ggd_jdict["packages"][ggd_recipe]["tags"]:
            if "uploaded_to_aws" in ggd_jdict["packages"][ggd_recipe]["tags"]["cached"]:
                print("\n\t-> The %s package is uploaded to an aws S3 bucket. To reduce processing time the package will be downloaded from an aws S3 bucket" %ggd_recipe)
                return(True)
    return(False)

def install_from_cached(ggd_recipe, ggd_channel,ggd_jdict,ggd_version,debug=False):
    """Method to install the ggd data package using a cached recipe

    install_from_cached
    ===================
    This method is used to install a ggd data package from a cached location. That is, a cached ggd recipe has
     been created and can be installed. Installing using a cached recipe increases the install speed. This is
     because (1) data processing and curation has already been done and the resulting files are cached. (This removes
     the time it takes to processes the data). (2) With a cached recipe we can bypass conda's solve environment step.

    If installed correctly the method returns True
    """

    conda_channel = "ggd-" + ggd_channel
    try:
        if debug:
            bypass_satsolver_on_install(ggd_recipe,conda_channel,debug=True)
        else:
            bypass_satsolver_on_install(ggd_recipe,conda_channel)

        get_file_locations(ggd_recipe,ggd_jdict,ggd_version)
        activate_enviroment_variables()
        print("\n\t-> DONE")

    except Exception as e:
        print("\n\t-> %s did not install properly. Review the error message:\n" %ggd_recipe)
        print(traceback.format_exc())
        check_for_installation(ggd_recipe,ggd_jdict) ## .uninstall method to remove extra ggd files
        print("\n\t-> %s was not installed. Please correct the errors and try again." %ggd_recipe)
        sys.exit(1)

    return(True)

def conda_install(ggd_recipe, ggd_channel,ggd_jdict,ggd_version,debug=False):
    """Method to install the recipe from the ggd-channel using conda

    conda_install
    ============
    This method is used to install the ggd recipe from the ggd conda channel using conda, if the files
     have not been cached.

    If installed correctly the method returns True
    """

    conda_version = get_required_conda_version()
    conda_install_str = "conda=" + conda_version
    if ggd_version != "-1":
        print("\n\t-> Installing %s version %s" %(ggd_recipe,ggd_version))
        try:
            if debug:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe+"="+ggd_version+"*", conda_install_str, "--debug"], stderr=sys.stderr, stdout=sys.stdout)
            else:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe+"="+ggd_version+"*", conda_install_str], stderr=sys.stderr, stdout=sys.stdout)
        except sp.CalledProcessError as e:
            sys.stderr.write("\n\t-> ERROR in install %s\n" %ggd_recipe)
            sys.stderr.write(str(e))
            sys.stderr.write("\n\t-> Removing files created by ggd during installation")
            check_for_installation(ggd_recipe,ggd_jdict) ## .uninstall method to remove extra ggd files
            sys.exit(e.returncode)
    else:
        print("\n\t-> Installing %s" %ggd_recipe)
        try:
            if debug:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe, conda_install_str, "--debug"], stderr=sys.stderr, stdout=sys.stdout)
            else:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe, conda_install_str], stderr=sys.stderr, stdout=sys.stdout)
        except sp.CalledProcessError as e:
            sys.stderr.write("\n\t-> ERROR in install %s\n" %ggd_recipe)
            sys.stderr.write(str(e))
            sys.stderr.write("\n\t-> Removing files created by ggd during installation")
            check_for_installation(ggd_recipe,ggd_jdict) ## .uninstall method to remove extra ggd files
            sys.exit(e.returncode)

    return(True)

def get_file_locations(ggd_recipe,ggd_jdict,ggd_version):
    """Method used to print the location of the installed files

    get_file_locations
    ==================
    This method is used to print the location of the data files installed for a reference
    for the user.
    """

    species = ggd_jdict["packages"][ggd_recipe]["identifiers"]["species"]
    build = ggd_jdict["packages"][ggd_recipe]["identifiers"]["genome-build"]
    version = ggd_jdict["packages"][ggd_recipe]["version"]
    CONDA_ROOT = conda_root()
    path = os.path.join(CONDA_ROOT,"share","ggd",species,build,ggd_recipe,version)
    print("\n\t-> Installation complete. The downloaded data files are located at:")
    print("\t\t%s" %path)
    print("\n\t-> A new environment variable that points to this directory path has also been created:")
    print("\t\t $ggd_%s\n" %ggd_recipe.replace("-","_"))

def install(parser, args):
    """Main method for installing a ggd data package

    install
    =======
    This method is the main method for running ggd install. It controls the different levels of install
    and file handeling.
    """

    print("\n\t-> Looking for %s in the 'ggd-%s' channel" %(args.name,args.channel))
    ## Check if the recipe is in ggd
    ggd_jsonDict = check_ggd_recipe(args.name,args.channel)
    if ggd_jsonDict == None:
        sys.exit()
    ## Check if the recipe is already installed
    if not check_if_installed(args.name,ggd_jsonDict,args.version):
        ## Check if conda has it installed on the system
        if not check_conda_installation(args.name,args.version):
            ## Check S3 bucket if version has not been set
            if args.version == "-1":
                if check_S3_bucket(args.name, ggd_jsonDict):
                    install_from_cached(args.name, args.channel, ggd_jsonDict, args.version, debug=args.debug)
                else:
                    if args.debug:
                        conda_install(args.name, args.channel, ggd_jsonDict,args.version,debug=True)
                    else:
                        conda_install(args.name, args.channel, ggd_jsonDict,args.version)

                    get_file_locations(args.name,ggd_jsonDict,args.version)
                    activate_enviroment_variables()
                    print("\n\t-> DONE")
            else:
                if args.debug:
                    conda_install(args.name, args.channel, ggd_jsonDict,args.version,debug=True)
                else:
                    conda_install(args.name, args.channel, ggd_jsonDict,args.version)

                get_file_locations(args.name,ggd_jsonDict,args.version)
                activate_enviroment_variables()
                print("\n\t-> DONE")

<commit_msg>Bug fixes<commit_after>#-------------------------------------------------------------------------------------------------------------
## Import Statements
#-------------------------------------------------------------------------------------------------------------
from __future__ import print_function
import sys
import os
import subprocess as sp
import glob
import traceback
from .utils import conda_root
from .utils import get_species
from .utils import get_ggd_channels
from .utils import get_channel_data
from .utils import get_channeldata_url
from .utils import bypass_satsolver_on_install
from .utils import active_conda_env
from .show_env import activate_enviroment_variables
from .search import load_json, load_json_from_url, search_packages
from .uninstall import remove_from_condaroot, check_for_installation
from .utils import get_required_conda_version

SPECIES_LIST = get_species()
#-------------------------------------------------------------------------------------------------------------
## Argument Parser
#-------------------------------------------------------------------------------------------------------------
def add_install(p):
    c = p.add_parser('install', help="install a data recipe from ggd")
    c.add_argument("-c", "--channel", default="genomics", choices=get_ggd_channels(),
                     help="The ggd channel the desired recipe is stored in. (Default = genomics)")
    c.add_argument("-v", "--version", default="-1", help="A specific ggd package version to install. If the -v flag is not used the latest version will be installed.")
    c.add_argument("-d", "--debug", action="store_true", help="(Optional) When the -d flag is set debug output will be printed to stdout.")
    c.add_argument("name", help="the name of the recipe to install")
    c.set_defaults(func=install)

#-------------------------------------------------------------------------------------------------------------
## Functions/Methods
#-------------------------------------------------------------------------------------------------------------

def check_ggd_recipe(ggd_recipe,ggd_channel):
    """Method used to check if the desired package is in the ggd repo using the repo metadata file

    check_ggd_recipe
    ================
    Method to check if the ggd recipe exists. Uses search_packages from search.py to
     search the ggd-channel json file. If the recipe exists within the json file,
     the installation proceeds. If not, the instalation stops
    """

    CHANNEL_DATA_URL = get_channeldata_url(ggd_channel)
    jdict = load_json_from_url(CHANNEL_DATA_URL)
    package_list = [x[0] for x in search_packages(jdict, ggd_recipe)]
    if ggd_recipe in package_list:
        print("\n\t-> %s exists in the ggd-%s channel" %(ggd_recipe,ggd_channel))
        return(jdict)
    else:
        print("\n\t-> '%s' was not found in ggd-%s" %(ggd_recipe, ggd_channel))
        print("\t-> You can search for recipes using the ggd search tool: \n\t\t'ggd search -t %s'\n" %ggd_recipe)
        return(None)

def check_if_installed(ggd_recipe,ggd_jdict,ggd_version):
    """Method to check if the recipe has already been installed and is in the conda ggd storage path.

    check_if_installed
    ==================
    This method is used to check if the ggd package has been installed and is located in the ggd storage path.
     If it is already installed the installation stops. If it is not detected then installation continues.
    """

    species = ggd_jdict["packages"][ggd_recipe]["identifiers"]["species"]
    build = ggd_jdict["packages"][ggd_recipe]["identifiers"]["genome-build"]
    version = ggd_jdict["packages"][ggd_recipe]["version"]

    CONDA_ROOT = conda_root()
    path = os.path.join(CONDA_ROOT,"share","ggd",species,build,ggd_recipe,version)
    recipe_exists = glob.glob(path)
    if recipe_exists:
        # if the ggd_version designated to be installed does not match the installed version, install the designated version
        if ggd_version != version and ggd_version != "-1":
            return(False)
        else:
            print("\n\t-> '%s' is already installed." %ggd_recipe)
            print("\t-> You can find %s here: %s" %(ggd_recipe,path))
            sys.exit()
    else:
        print("\n\t-> %s version %s is not installed on your system" %(ggd_recipe, version))
        return(False)

def check_conda_installation(ggd_recipe,ggd_version):
    """Method used to check if the recipe has been installed using conda.

    check_conda_installation
    ========================
    This method is used to check if the ggd data package has been installed by the conda system,
     without being installed by the ggd system. If it has, the recipe needs to be uninstalled and
     reinstalled. If not, the system continues to install the package using ggd.
    """

    conda_package_list = sp.check_output(["conda", "list"]).decode('utf8')
    recipe_find = conda_package_list.find(ggd_recipe)
    if recipe_find == -1:
        print("\n\t-> %s has not been installed by conda" %ggd_recipe)
        return(False)
    else:
        if ggd_version != "-1": ## Check if ggd version was designated
            installed_version = conda_package_list[recipe_find:recipe_find+100].split("\n")[0].replace(" ","")[len(ggd_recipe)]
            if installed_version != ggd_version:
                print("\n\t-> %s version %s has not been installed by conda" %(ggd_recipe,str(ggd_version)))
                return(False)
            else:
                print("\n\t-> %s version %s has been installed by conda on your system and must be uninstalled to proceed." %(ggd_recipe,str(ggd_version)))
                print("\t-> To reinstall run:\n\t\t $ ggd uninstall %s \n\t\t $ ggd install %s" %(ggd_recipe,ggd_recipe))
                sys.exit()
        else: ## If version is not specified check if exact package in conda list
            start_index = conda_package_list.find(ggd_recipe)
            end_index = start_index + len(ggd_recipe)
            ## Check if it really is the package, or something similar to the package
            if conda_package_list[int(start_index) - 1] == "\n" and conda_package_list[int(end_index) + 1] == " ":
                print("\n\t-> %s has been installed by conda on your system and must be uninstalled to proceed." %ggd_recipe)
                print("\t-> To reinstall run:\n\t\t $ ggd uninstall %s \n\t\t $ ggd install %s" %(ggd_recipe,ggd_recipe))
                sys.exit()
            else: ## IF not exactly in conda list
                print("\n\t-> %s has not been installed by conda" %ggd_recipe)
                return(False)

def check_S3_bucket(ggd_recipe, ggd_jdict):
    """Method to check if the recipe is stored on the ggd S3 bucket. If so it installs from S3

    check_S3_bucket
    ==============
    This method is used to check if the recipe has been cached on aws S3 bucket. It returns true if the
     the recipe is cached, and false if it is not. If it is cached the cached version will be installed.
    """

    if "tags" in ggd_jdict["packages"][ggd_recipe]:
        if "cached" in ggd_jdict["packages"][ggd_recipe]["tags"]:
            if "uploaded_to_aws" in ggd_jdict["packages"][ggd_recipe]["tags"]["cached"]:
                print("\n\t-> The %s package is uploaded to an aws S3 bucket. To reduce processing time the package will be downloaded from an aws S3 bucket" %ggd_recipe)
                return(True)
    return(False)

def install_from_cached(ggd_recipe, ggd_channel,ggd_jdict,ggd_version,debug=False):
    """Method to install the ggd data package using a cached recipe

    install_from_cached
    ===================
    This method is used to install a ggd data package from a cached location. That is, a cached ggd recipe has
     been created and can be installed. Installing using a cached recipe increases the install speed. This is
     because (1) data processing and curation has already been done and the resulting files are cached. (This removes
     the time it takes to processes the data). (2) With a cached recipe we can bypass conda's solve environment step.

    If installed correctly the method returns True
    """

    conda_channel = "ggd-" + ggd_channel
    try:
        if debug:
            bypass_satsolver_on_install(ggd_recipe,conda_channel,debug=True)
        else:
            bypass_satsolver_on_install(ggd_recipe,conda_channel)

        get_file_locations(ggd_recipe,ggd_jdict,ggd_version)
        activate_enviroment_variables()
        print("\n\t-> DONE")

    except Exception as e:
        print("\n\t-> %s did not install properly. Review the error message:\n" %ggd_recipe)
        print(traceback.format_exc())
        check_for_installation(ggd_recipe,ggd_jdict) ## .uninstall method to remove extra ggd files
        print("\n\t-> %s was not installed. Please correct the errors and try again." %ggd_recipe)
        sys.exit(1)

    return(True)

def conda_install(ggd_recipe, ggd_channel,ggd_jdict,ggd_version,debug=False):
    """Method to install the recipe from the ggd-channel using conda

    conda_install
    ============
    This method is used to install the ggd recipe from the ggd conda channel using conda, if the files
     have not been cached.

    If installed correctly the method returns True
    """

    conda_version = get_required_conda_version()
    conda_install_str = "conda=" + conda_version
    if ggd_version != "-1":
        print("\n\t-> Installing %s version %s" %(ggd_recipe,ggd_version))
        try:
            if debug:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe+"="+str(ggd_version)+"*", conda_install_str, "--debug"], stderr=sys.stderr, stdout=sys.stdout)
            else:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe+"="+str(ggd_version)+"*", conda_install_str], stderr=sys.stderr, stdout=sys.stdout)
        except sp.CalledProcessError as e:
            sys.stderr.write("\n\t-> ERROR in install %s\n" %ggd_recipe)
            sys.stderr.write(str(e))
            sys.stderr.write("\n\t-> Removing files created by ggd during installation")
            check_for_installation(ggd_recipe,ggd_jdict) ## .uninstall method to remove extra ggd files
            sys.exit(e.returncode)
    else:
        print("\n\t-> Installing %s" %ggd_recipe)
        try:
            if debug:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe, conda_install_str, "--debug"], stderr=sys.stderr, stdout=sys.stdout)
            else:
                sp.check_call(["conda", "install", "-c", "ggd-"+ggd_channel, "-y", ggd_recipe, conda_install_str], stderr=sys.stderr, stdout=sys.stdout)
        except sp.CalledProcessError as e:
            sys.stderr.write("\n\t-> ERROR in install %s\n" %ggd_recipe)
            sys.stderr.write(str(e))
            sys.stderr.write("\n\t-> Removing files created by ggd during installation")
            check_for_installation(ggd_recipe,ggd_jdict) ## .uninstall method to remove extra ggd files
            sys.exit(e.returncode)

    return(True)

def get_file_locations(ggd_recipe,ggd_jdict,ggd_version):
    """Method used to print the location of the installed files

    get_file_locations
    ==================
    This method is used to print the location of the data files installed for a reference
    for the user.
    """

    species = ggd_jdict["packages"][ggd_recipe]["identifiers"]["species"]
    build = ggd_jdict["packages"][ggd_recipe]["identifiers"]["genome-build"]
    version = ggd_jdict["packages"][ggd_recipe]["version"]
    CONDA_ROOT = conda_root()
    path = os.path.join(CONDA_ROOT,"share","ggd",species,build,ggd_recipe,version)
    print("\n\t-> Installation complete. The downloaded data files are located at:")
    print("\t\t%s" %path)
    print("\n\t-> A new environment variable that points to this directory path has also been created:")
    print("\t\t $ggd_%s\n" %ggd_recipe.replace("-","_"))

def install(parser, args):
    """Main method for installing a ggd data package

    install
    =======
    This method is the main method for running ggd install. It controls the different levels of install
    and file handeling.
    """

    print("\n\t-> Looking for %s in the 'ggd-%s' channel" %(args.name,args.channel))
    ## Check if the recipe is in ggd
    ggd_jsonDict = check_ggd_recipe(args.name,args.channel)
    if ggd_jsonDict == None:
        sys.exit()
    ## Check if the recipe is already installed
    if not check_if_installed(args.name,ggd_jsonDict,args.version):
        ## Check if conda has it installed on the system
        if not check_conda_installation(args.name,args.version):
            ## Check S3 bucket if version has not been set
            if args.version == "-1":
                if check_S3_bucket(args.name, ggd_jsonDict):
                    install_from_cached(args.name, args.channel, ggd_jsonDict, args.version, debug=args.debug)
                else:
                    if args.debug:
                        conda_install(args.name, args.channel, ggd_jsonDict,args.version,debug=True)
                    else:
                        conda_install(args.name, args.channel, ggd_jsonDict,args.version)

                    get_file_locations(args.name,ggd_jsonDict,args.version)
                    activate_enviroment_variables()
                    print("\n\t-> DONE")
            else:
                if args.debug:
                    conda_install(args.name, args.channel, ggd_jsonDict,args.version,debug=True)
                else:
                    conda_install(args.name, args.channel, ggd_jsonDict,args.version)

                get_file_locations(args.name,ggd_jsonDict,args.version)
                activate_enviroment_variables()
                print("\n\t-> DONE")
    return(True)
<|endoftext|>


<commit_before>from django.shortcuts import render, redirect
from django.utils import timezone, translation
from django.core.mail import send_mail, BadHeaderError
from apps.core.models import Notice, Statistic, Service
from django.core.exceptions import SuspiciousOperation
from apps.core.backends import validate_recaptcha
import logging
import json

logger = logging.getLogger('sso')

# /main/
def main(request):
    current_time = timezone.now()
    services = Service.objects.filter(is_shown=True, icon__isnull=False)
    notice = Notice.objects.filter(valid_from__lte=current_time,
                                   valid_to__gt=current_time).first()

    return render(request, 'main.html', {'services': services, 'notice': notice})

# /lang/
def lang(request, code):
    if code not in ['en', 'ko']:
        return redirect('/')

    translation.activate(code)
    request.session[translation.LANGUAGE_SESSION_KEY] = code
    return redirect(request.META.get('HTTP_REFERER', '/'))

# /credits/
def credits(request):
    return render(request, 'credits.html')

# /terms/
def terms(request):
    return render(request, 'terms.html')

# /privacy/
def privacy(request):
    return render(request, 'privacy.html')

# /stats/
def stats(request):
    level = 0
    if request.user.is_authenticated():
        if request.user.is_staff:
            level = 2
        elif request.user.profile.sparcs_id:
            level = 1

    raw_stats = Statistic.objects.order_by('-time')
    if len(raw_stats) == 0:
        raw_stat = {}
    else:
        time = raw_stats[0].time
        raw_stat = json.loads(raw_stats[0].data)

    stat = []
    for name, value in raw_stat.iteritems():
        if name != 'all':
            service = Service.objects.get(name=name)
            if level < 1 and not service.is_shown:
                continue

        s = {}
        s['name'] = name
        s['alias'] = service.alias if name != 'all' else 'all'
        s.update(value)

        if name == 'all':
            stat.insert(0, s)
        else:
            stat.append(s)

    return render(request, 'stats.html', {'level': level, 'time': time, 'stat': stat})

# /help/
def help(request):
    return render(request, 'help.html')

# /contact/
def contact(request):
    name = ''
    email = ''
    subject = ''
    message = ''
    if request.method == 'POST':
        subject = request.POST.get('subject','')
        name = request.POST.get('name','')
        email = request.POST.get('email','')
        message = request.POST.get('message','')
        recipients = ['gogi@sparcs.org']
        real_subject = "Sparcs SSO Report : " + subject

        result = validate_recaptcha(request.POST.get('g-recaptcha-response',''))

        if not result:
            return redirect('/')

        if subject and name and email and message:
            try:
                contents = 'subject : ' + subject + '\nname : ' + name + '\nmessage\n' + message
                send_mail(real_subject, contents, email, recipients)
            except ValueError:
                raise SuspiciousOperation()
            return redirect('/thanks/')
    return render(request, 'contact.html',
            {'name': name, 'email': email, 'subject': subject, 'message':message})

# /thanks/
def thanks(request):
    state = ''
    if request.method == 'POST':
        state = requset.POST.get('state','')
        if state:
            return redirect('/')
    return render(request, 'thanks.html')
<commit_msg>for merge<commit_after>from django.shortcuts import render, redirect
from django.utils import timezone, translation
from django.core.mail import send_mail, BadHeaderError
from apps.core.models import Notice, Statistic, Service
from django.core.exceptions import SuspiciousOperation
from apps.core.backends import validate_recaptcha
import logging
import json

logger = logging.getLogger('sso')

# /main/
def main(request):
    current_time = timezone.now()
    services = Service.objects.filter(is_shown=True, icon__isnull=False)
    notice = Notice.objects.filter(valid_from__lte=current_time,
                                   valid_to__gt=current_time).first()

    return render(request, 'main.html', {'services': services, 'notice': notice})

# /lang/
def lang(request, code):
    if code not in ['en', 'ko']:
        return redirect('/')

    translation.activate(code)
    request.session[translation.LANGUAGE_SESSION_KEY] = code
    return redirect(request.META.get('HTTP_REFERER', '/'))

# /credits/
def credits(request):
    return render(request, 'credits.html')

# /terms/
def terms(request):
    return render(request, 'terms.html')

# /privacy/
def privacy(request):
    return render(request, 'privacy.html')

# /stats/
def stats(request):
    level = 0
    if request.user.is_authenticated():
        if request.user.is_staff:
            level = 2
        elif request.user.profile.sparcs_id:
            level = 1

    raw_stats = Statistic.objects.order_by('-time')
    if len(raw_stats) == 0:
        raw_stat = {}
    else:
        time = raw_stats[0].time
        raw_stat = json.loads(raw_stats[0].data)

    stat = []
    for name, value in raw_stat.iteritems():
        if name != 'all':
            service = Service.objects.get(name=name)
            if level < 1 and not service.is_shown:
                continue

        s = {}
        s['name'] = name
        s['alias'] = service.alias if name != 'all' else 'all'
        s.update(value)

        if name == 'all':
            stat.insert(0, s)
        else:
            stat.append(s)

    return render(request, 'stats.html', {'level': level, 'time': time, 'stat': stat})

# /help/
def help(request):
    return render(request, 'help.html')

# /contact/
def contact(request):
    if request.method == 'POST':
        subject = request.POST.get('subject','')
        name = request.POST.get('name','')
        email = request.POST.get('email','')
        message = request.POST.get('message','')
        recipients = ['gogi@sparcs.org']
        real_subject = "SPARCS SSO Report : " + subject

        result = validate_recaptcha(request.POST.get('g-recaptcha-response',''))

        if not result:
            return redirect('/')

        if subject and name and email and message:
            try:
                contents = 'subject : ' + subject + '\nname : ' + name + '\nmessage\n' + message
                send_mail(real_subject, contents, email, recipients)
            except ValueError:
                raise SuspiciousOperation()
            return redirect('/thanks/')
    return render(request, 'contact.html')

# /thanks/
def thanks(request):
    state = ''
    if request.method == 'POST':
        state = requset.POST.get('state','')
        if state:
            return redirect('/')
    return render(request, 'thanks.html')
<|endoftext|>


<commit_before>import socket
import mimetypes
from os import getcwd
from os.path import isdir
from os.path import isfile
from os import listdir

class InvalidHttpCodeError(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class NotGETRequestError(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class NotHTTP1_1Error(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class BadRequestError(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class ResourceNotFound(Exception):
    u"""An exception that is raised when a resources cannot be located
          for the provided URI"""
    pass

class HttpServer(object):
    u"""Create an HTTP Server with the given endpoint."""
    def __init__(self, ip=u'127.0.0.1', port=50000, backlog=5):
        self._ip = ip
        self._port = port
        self._backlog = backlog
        self._socket = None
        self._root = getcwd() + b"/root/"
        self._statusCodes = {
            200: b'OK',
            301: b'Moved Permanently',
            304: b'Not Modified',
            400: b'Bad Request',
            404: b'Not Found',
            405: b'Method Not Allowed',
            500: b'Internal Server Error',
            505: b'HTTP version not supported'}

    def open_socket(self):
        u"""Open a socket, bind it, and listen."""
        self._socket = socket.socket(
            socket.AF_INET,
            socket.SOCK_STREAM,
            socket.IPPROTO_IP)
        self._socket.bind((self._ip, self._port))
        self._socket.listen(self._backlog)

    def close_socket(self):
        u"""Shutdown and close the socket."""
        self._socket.close()
        self._socket = None

    def gen_response(self, code, body=None, kwargs={}):
        u"""Generate response for the given HTTP status code."""
        response = []
        try:
            response.append("HTTP/1.1 {} {}\r\n".format(
                code, self._statusCodes[code]))
            headers = kwargs.keys()
            for header in headers:
                    response.append("{}: {}\r\n".format(
                        header,
                        kwargs[header]))
            response.append("\r\n")  # blank line between headers and body
            if body:
                response.append(body)
            return "".join(response)
        except KeyError:
            raise InvalidHttpCodeError(
                u'{} is not a valid HTTP code'.format(code))

    def start_listening(self):
        u"""Accept a request from a client and return appropriate response."""
        buffersize = 32
        while True:
            request = []
            connection, addr = self._socket.accept()
            while True:
                buffer_ = connection.recv(buffersize)
                request.append(buffer_)
                if len(buffer_) < buffersize:
                    break
            try:
                body, content_type = self.process_request("".join(request))
            except NotGETRequestError:
                connection.sendall(self.gen_response(405))
            except BadRequestError:
                connection.sendall(self.gen_response(400))
            except NotHTTP1_1Error:
                connection.sendall(self.gen_response(505))
            except ResourceNotFound:
                connection.sendall(self.gen_response(404))
            else:
                response = self.gen_response(
                    200,
                    body,
                    {'Content-Type': content_type,
                     'Content-Length': len(body)})
                connection.sendall(response)
                print response
            finally:
                connection.close()

    def process_request(self, request):
        u"""Split status line into verb, URI, and protocol."""
        list_ = request.split("\r\n")
        status_line = list_[0].split()
        if len(status_line) != 3:
            raise BadRequestError
        if status_line[0] != "GET":
            raise NotGETRequestError
        if status_line[2] != "HTTP/1.1":
            raise NotHTTP1_1Error
        body, content_type = self._retrieve_resource(status_line[1])
        return body, content_type

    def _retrieve_resource(self, uri):
        u"""
        Get the resource specified by the uri if it exist.
        Otherwise, raise a Exception
        """
        p = self._root + uri
        if isdir(p):
            body = ["
Directory Listing for "]
            body.append(uri)
            body.append("

")
            for res in listdir(p):
                body.append("
 {}
".format(res))
            body.append("
")
            return ("".join(body), "text/html")
        elif isfile(p):
            with open(self._root + uri, 'r') as resource:
                body = resource.read()
                content_type, content_encoding = mimetypes.guess_type(uri)
            return (body, content_type)
        else:
            raise ResourceNotFound

if __name__ == "__main__":
    s = HttpServer()
    try:
        s.open_socket()
        s.start_listening()
    except KeyboardInterrupt:
        s.close_socket()
    finally:
        if s is not None and s._socket is not None:
            s.close_socket()
<commit_msg>Revise HttpServer._retrieve_resource to display links to resources in a directory and to append / to sub directory URIs<commit_after>import socket
import mimetypes
from os import getcwd
from os.path import isdir
from os.path import isfile
from os import listdir

class InvalidHttpCodeError(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class NotGETRequestError(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class NotHTTP1_1Error(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class BadRequestError(Exception):
    u"""Create InvalidHttpCodeError Exception"""
    pass

class ResourceNotFound(Exception):
    u"""An exception that is raised when a resources cannot be located
          for the provided URI"""
    pass

class HttpServer(object):
    u"""Create an HTTP Server with the given endpoint."""
    def __init__(self, ip=u'127.0.0.1', port=50000, backlog=5):
        self._ip = ip
        self._port = port
        self._backlog = backlog
        self._socket = None
        self._root = getcwd() + b"/root/"
        self._statusCodes = {
            200: b'OK',
            301: b'Moved Permanently',
            304: b'Not Modified',
            400: b'Bad Request',
            404: b'Not Found',
            405: b'Method Not Allowed',
            500: b'Internal Server Error',
            505: b'HTTP version not supported'}

    def open_socket(self):
        u"""Open a socket, bind it, and listen."""
        self._socket = socket.socket(
            socket.AF_INET,
            socket.SOCK_STREAM,
            socket.IPPROTO_IP)
        self._socket.bind((self._ip, self._port))
        self._socket.listen(self._backlog)

    def close_socket(self):
        u"""Shutdown and close the socket."""
        self._socket.close()
        self._socket = None

    def gen_response(self, code, body=None, kwargs={}):
        u"""Generate response for the given HTTP status code."""
        response = []
        try:
            response.append("HTTP/1.1 {} {}\r\n".format(
                code, self._statusCodes[code]))
            headers = kwargs.keys()
            for header in headers:
                    response.append("{}: {}\r\n".format(
                        header,
                        kwargs[header]))
            response.append("\r\n")  # blank line between headers and body
            if body:
                response.append(body)
            return "".join(response)
        except KeyError:
            raise InvalidHttpCodeError(
                u'{} is not a valid HTTP code'.format(code))

    def start_listening(self):
        u"""Accept a request from a client and return appropriate response."""
        buffersize = 32
        while True:
            request = []
            connection, addr = self._socket.accept()
            while True:
                buffer_ = connection.recv(buffersize)
                request.append(buffer_)
                if len(buffer_) < buffersize:
                    break
            try:
                body, content_type = self.process_request("".join(request))
            except NotGETRequestError:
                connection.sendall(self.gen_response(405))
            except BadRequestError:
                connection.sendall(self.gen_response(400))
            except NotHTTP1_1Error:
                connection.sendall(self.gen_response(505))
            except ResourceNotFound:
                connection.sendall(self.gen_response(404))
            else:
                response = self.gen_response(
                    200,
                    body,
                    {'Content-Type': content_type,
                     'Content-Length': len(body)})
                connection.sendall(response)
                print response
            finally:
                connection.close()

    def process_request(self, request):
        u"""Split status line into verb, URI, and protocol."""
        list_ = request.split("\r\n")
        status_line = list_[0].split()
        if len(status_line) != 3:
            raise BadRequestError
        if status_line[0] != "GET":
            raise NotGETRequestError
        if status_line[2] != "HTTP/1.1":
            raise NotHTTP1_1Error
        body, content_type = self._retrieve_resource(status_line[1])
        return body, content_type

    def _retrieve_resource(self, uri):
        u"""
        Get the resource specified by the uri if it exist.
        Otherwise, raise a Exception
        """
        p = self._root + uri
        if isdir(p):
            body = ["
Directory Listing for "]
            body.append(uri)
            body.append("

")
            for res in listdir(p):
                if isdir(p+res):
                    res += b'/'
                body.append('
{}
'.format(res, res))
            body.append("
")
            return ("".join(body), "text/html")
        elif isfile(p):
            with open(self._root + uri, 'r') as resource:
                body = resource.read()
                content_type, content_encoding = mimetypes.guess_type(uri)
            return (body, content_type)
        else:
            raise ResourceNotFound

if __name__ == "__main__":
    s = HttpServer()
    try:
        s.open_socket()
        s.start_listening()
    except KeyboardInterrupt:
        s.close_socket()
    finally:
        if s is not None and s._socket is not None:
            s.close_socket()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from openerp import models, fields, api
from datetime import datetime, timedelta
from openerp.tools import DEFAULT_SERVER_DATETIME_FORMAT

class AccountAnalyticAccountUsedSlots(models.Model):
    _inherit = 'account.analytic.account'

    used_slots = fields.Integer(default=0, help='number of used slots', readonly=True)
    remind_on_slots = fields.Integer(help='configure when to remind a customer about remaining slots', string='Remind on (slots)', default=2)

    @api.model
    def _cron_compute_used_slots(self):
        for record in self.search([]):
            lines = self.env['sale.order.line'].search([
                ('order_id.project_id', '=', record.id),
                ('booking_end', '!=', False),
                ('booking_end', '<=', datetime.now().strftime(DEFAULT_SERVER_DATETIME_FORMAT)),
                ('order_id.state', 'in', ['manual', 'done']),
                ('price_unit', '=', '0'),
            ])
            uslots = 0
            for line in lines:
                uslots += line.product_uom_qty

            record.used_slots = uslots

class SaleOrderTheCage(models.Model):
    _inherit = 'sale.order'

    expiring_reminder = fields.Boolean(default=False)

    @api.one
    def write(self, vals):
        result = super(SaleOrderTheCage, self).write(vals)
        if vals.get('state') == 'sent':
            msg = 'Sale Order #' + self.name + ' is confirmed'
            phone = self.partner_id.mobile
            self.env['sms_sg.sendandlog'].send_sms(phone, msg)
        return result

class SaleOrderLineReminder(models.Model):
    _inherit = 'sale.order.line'

    booking_reminder = fields.Boolean(default=False, select=True)

    @api.model
    def _cron_booking_reminder(self):
        lines = self.search(['&', '&', '&', ('booking_reminder', '=', False),
                             ('booking_start', '!=', False),
                             ('booking_start', '<=', (datetime.now() + timedelta(hours=48)).strftime(DEFAULT_SERVER_DATETIME_FORMAT)),
                             '|',
                             ('order_id.state', '=', 'done'),
                             ('price_unit', '=', '0'),
                             ])
        lines.write({'booking_reminder': True})
        for line in lines:
            msg = 'Sale Order #' + line.order_id.name + ' is confirmed'
            phone = line.order_id.partner_id.mobile
            self.env['sms_sg.sendandlog'].send_sms(phone, msg)

class ResPartnerReminderConfig(models.Model):
    _inherit = 'res.partner'

    reminder_sms = fields.Boolean(default=True, string='Booking sms reminder enabled')
    reminder_email = fields.Boolean(default=True, string='Booking email reminder enabled')
<commit_msg>[ADD] reminder_sms usage when sending sms<commit_after># -*- coding: utf-8 -*-
from openerp import models, fields, api
from datetime import datetime, timedelta
from openerp.tools import DEFAULT_SERVER_DATETIME_FORMAT

class AccountAnalyticAccountUsedSlots(models.Model):
    _inherit = 'account.analytic.account'

    used_slots = fields.Integer(default=0, help='number of used slots', readonly=True)
    remind_on_slots = fields.Integer(help='configure when to remind a customer about remaining slots', string='Remind on (slots)', default=2)

    @api.model
    def _cron_compute_used_slots(self):
        for record in self.search([]):
            lines = self.env['sale.order.line'].search([
                ('order_id.project_id', '=', record.id),
                ('booking_end', '!=', False),
                ('booking_end', '<=', datetime.now().strftime(DEFAULT_SERVER_DATETIME_FORMAT)),
                ('order_id.state', 'in', ['manual', 'done']),
                ('price_unit', '=', '0'),
            ])
            uslots = 0
            for line in lines:
                uslots += line.product_uom_qty

            record.used_slots = uslots

class SaleOrderTheCage(models.Model):
    _inherit = 'sale.order'

    expiring_reminder = fields.Boolean(default=False)

    @api.one
    def write(self, vals):
        result = super(SaleOrderTheCage, self).write(vals)
        if vals.get('state') == 'sent':  # send sms immediately after user pushed 'Send by Email' button on the Sale Order
            if self.partner_id.reminder_sms:
                msg = 'Sale Order #' + self.name + ' is confirmed'
                phone = self.partner_id.mobile
                self.env['sms_sg.sendandlog'].send_sms(phone, msg)
        return result

class SaleOrderLineReminder(models.Model):
    _inherit = 'sale.order.line'

    booking_reminder = fields.Boolean(default=False, select=True)

    @api.model
    def _cron_booking_reminder(self):
        lines = self.search(['&', '&', '&', ('booking_reminder', '=', False),
                             ('booking_start', '!=', False),
                             ('booking_start', '<=', (datetime.now() + timedelta(hours=48)).strftime(DEFAULT_SERVER_DATETIME_FORMAT)),
                             '|',
                             ('order_id.state', '=', 'done'),
                             ('price_unit', '=', '0'),
                             ])
        lines.write({'booking_reminder': True})
        for line in lines:
            msg = 'Sale Order #' + line.order_id.name + ' is confirmed'
            phone = line.order_id.partner_id.mobile
            self.env['sms_sg.sendandlog'].send_sms(phone, msg)

class ResPartnerReminderConfig(models.Model):
    _inherit = 'res.partner'

    reminder_sms = fields.Boolean(default=True, string='Booking sms reminder enabled')
    reminder_email = fields.Boolean(default=True, string='Booking email reminder enabled')
<|endoftext|>


<commit_before>"""
The Message class represents a message that is sent or received and contains
methods for publishing the message, or in the case that the message was
delivered by RabbitMQ, acknowledging it, rejecting it or negatively
acknowledging it.

"""
import datetime
import json
import logging
import math
import time
import pprint
import uuid

from pamqp import body
from pamqp import header
from pamqp import specification

from rabbitpy import DEBUG
from rabbitpy import base
from rabbitpy import exceptions

LOGGER = logging.getLogger(__name__)

class Properties(specification.Basic.Properties):
    """Proxy class for :py:class:`pamqp.specification.Basic.Properties`"""
    pass

class Message(base.AMQPClass):
    """Created by both rabbitpy internally when a message is delivered or
    returned from RabbitMQ and by implementing applications, the Message class
    is used to publish a message to and access and respond to a message from
    RabbitMQ.

    When specifying properties for a message, pass in a dict of key value items
    that match the AMQP Basic.Properties specification with a small caveat.

    Due to an overlap in the AMQP specification and the Python keyword
    :code:`type`, the :code:`type` property is referred to as
    :code:`message_type`.

    The following is a list of the available properties:

    * app_id
    * content_type
    * content_encoding
    * correlation_id
    * delivery_node
    * expiration
    * headers
    * message_id
    * message_type
    * priority
    * reply_to
    * timestamp
    * user_id

    **Automated features**

    When passing in the body value, if it is a dict or list, it will
    automatically be JSON serialized and the content type "application/json"
    will be set on the message properties.

    When publishing a message to RabbitMQ, if the auto_id value is True and no
    message_id value was passed in as a property, a UUID will be generated and
    specified as a property of the message.

    If a timestamp is not specified when passing in properties, the current
    Unix epoch value will be set in the message properties.

    :param channel: The channel object for the message object to act upon
    :type channel: :py:class:`rabbitpy.channel.Channel`
    :param str or dict or list body_value: The message body
    :param dict properties: A dictionary of message properties
    :param bool auto_id: Add a message id if no properties were passed in.

    """
    method = None
    name = 'Message'

    def __init__(self, channel, body_value, properties=None, auto_id=True):
        """Create a new instance of the Message object."""
        super(Message, self).__init__(channel, 'Message')

        if isinstance(body_value, dict) or isinstance(body_value, list):
            body_value = json.dumps(body_value, ensure_ascii=False)
            if properties is None:
                properties = {'content_type': 'application/json'}
            else:
                properties['content_type'] = 'application/json'
        self.body = body_value
        self.properties = properties or self._base_properties
        if auto_id and 'message_id' not in self.properties:
            self._add_auto_message_id()
        if 'timestamp' not in self.properties:
            self._add_timestamp()
        if isinstance(self.properties['timestamp'], time.struct_time):
            self.properties['timestamp'] = self._timestamp_from_struct_time()

    def _add_auto_message_id(self):
        """Set the message_id property to a new UUID."""
        self.properties['message_id'] = str(uuid.uuid4())

    def _add_timestamp(self):
        """Add the timestamp to the properties"""
        self.properties['timestamp'] = datetime.datetime.now()

    def _timestamp_from_struct_time(self):
        return datetime.datetime(*self.properties['timestamp'][:6])

    @property
    def _base_properties(self):
        """Return a base set of properties if no properties were passed into
        the constructor.

        :rtype: dict

        """
        return {"message_id": str(uuid.uuid4()),
                "timestamp": datetime.datetime.now()}

    @property
    def _properties(self):
        """Return a new Basic.Properties object representing the message
        properties.

        :rtype: pamqp.specification.Basic.Properties

        """
        invalid_keys = [key for key in self.properties
                        if key not in
                        specification.Basic.Properties.attributes]
        self._prune_invalid_properties(invalid_keys)
        self._coerce_properties()
        return specification.Basic.Properties(**self.properties)

    def _coerce_properties(self):
        """Force properties to be set to the correct data type"""
        for key in self.properties:
            _type = getattr(specification.Basic.Properties, key)
            if DEBUG:
                LOGGER.debug('Type: %s, %s', _type, type(self.properties[key]))
            if _type == 'shortstr' and not isinstance(self.properties[key],
                                                      basestring):
                LOGGER.warning('Coercing property %s to bytes', key)
                self.properties[key] = bytes(self.properties[key])
            elif _type == 'octet' and not isinstance(self.properties[key],
                                                     int):
                LOGGER.warning('Coercing property %s to int', key)
                self.properties[key] = int(self.properties[key])
            elif _type == 'table' and not isinstance(self.properties[key],
                                                     dict):
                LOGGER.warning('Resetting invalid value for %s to None', key)
                self.properties[key] = {}
            elif _type == 'timestamp' and isinstance(self.properties[key],
                                                     int):
                LOGGER.warning('Coercing property %s to datetime', key)
                self.properties[key] = \
                    datetime.datetime.fromtimestamp(self.properties[key])

    def _prune_invalid_properties(self, invalid_keys):
        """Remove invalid properties from the message properties.

        :param list invalid_keys: A list of invalid property names to remove

        """
        for key in invalid_keys:
            LOGGER.warning('Removing invalid property "%s"', key)
            del self.properties[key]

    def ack(self, all_previous=False):
        """Acknowledge receipt of the message to RabbitMQ. Will raise an
        ActionException if the message was not received from a broker.

        :raises: ActionException

        """
        if not self.method:
            raise exceptions.ActionException('Can not ack non-received '
                                             'message')
        basic_ack = specification.Basic.Ack(self.method.delivery_tag,
                                            multiple=all_previous)
        self.channel._write_frame(basic_ack)

    def json(self):
        """Deserialize the message body if it is JSON, returning the value.

        :rtype: any

        """
        return json.loads(self.body)

    def nack(self, requeue=False, all_previous=False):
        """Negatively acknowledge receipt of the message to RabbitMQ. Will
        raise an ActionException if the message was not received from a broker.

        :param bool requeue: Requeue the message
        :param bool all_previous: Nack all previous unacked messages up to and
                                  including this one
        :raises: ActionException

        """
        if not self.method:
            raise exceptions.ActionException('Can not nack non-received '
                                             'message')
        basic_nack = specification.Basic.Nack(self.method.delivery_tag,
                                              requeue=requeue,
                                              multiple=all_previous)
        self.channel._write_frame(basic_nack)

    def pprint(self, properties=False):
        """Print a formatted representation of the message.

        :param bool properties: Include properties in the representation

        """
        print('Exchange: %s\n' % self.method.exchange)
        print('Routing Key: %s\n' % self.method.routing_key)
        if properties:
            print('Properties:\n')
            pprint.pprint(self.properties)
            print('\nBody:\n')
        pprint.pprint(self.body)

    def publish(self, exchange, routing_key='', mandatory=False):
        """Publish the message to the exchange with the specified routing
        key.

        :param str | rabbitpy.base.AMQPClass exchange: The exchange to bind to
        :param str routing_key: The routing key to use
        :param bool mandatory: Requires the message is published
        :return: bool | None
        :raises: rabbitpy.exceptions.MessageReturnedException

        """
        if isinstance(exchange, base.AMQPClass):
            exchange = exchange.name
        method_frame = specification.Basic.Publish(exchange=exchange,
                                                   routing_key=
                                                   routing_key or '',
                                                   mandatory=mandatory)
        self.channel._write_frame(method_frame)
        header_frame = header.ContentHeader(body_size=len(self.body),
                                            properties=self._properties)
        self.channel._write_frame(header_frame)
        pieces = int(math.ceil(len(self.body) /
                               float(self.channel.maximum_frame_size)))

        # Send the message
        for offset in range(0, pieces):
            start = self.channel.maximum_frame_size * offset
            end = start + self.channel.maximum_frame_size
            if end > len(self.body):
                end = len(self.body)
            self.channel._write_frame(body.ContentBody(self.body[start:end]))

        # If publisher confirmations are enabled, wait for the response
        if self.channel.publisher_confirms:
            response = self.channel._wait_for_confirmation()
            if isinstance(response, specification.Basic.Ack):
                return True
            elif isinstance(response, specification.Basic.Nack):
                return False
            else:
                raise exceptions.UnexpectedResponseError(response)

    def reject(self, requeue=False):
        """Reject receipt of the message to RabbitMQ. Will raise
        an ActionException if the message was not received from a broker.

        :param bool requeue: Requeue the message
        :raises: ActionException

        """
        if not self.method:
            raise exceptions.ActionException('Can not reject non-received '
                                             'message')
        basic_reject = specification.Basic.Reject(self.method.delivery_tag,
                                                  requeue=requeue)
        self.channel._write_frame(basic_reject)
<commit_msg>Add delivery_tag as a read-only property of rabbitpy.Message<commit_after>"""
The Message class represents a message that is sent or received and contains
methods for publishing the message, or in the case that the message was
delivered by RabbitMQ, acknowledging it, rejecting it or negatively
acknowledging it.

"""
import datetime
import json
import logging
import math
import time
import pprint
import uuid

from pamqp import body
from pamqp import header
from pamqp import specification

from rabbitpy import DEBUG
from rabbitpy import base
from rabbitpy import exceptions

LOGGER = logging.getLogger(__name__)

class Properties(specification.Basic.Properties):
    """Proxy class for :py:class:`pamqp.specification.Basic.Properties`"""
    pass

class Message(base.AMQPClass):
    """Created by both rabbitpy internally when a message is delivered or
    returned from RabbitMQ and by implementing applications, the Message class
    is used to publish a message to and access and respond to a message from
    RabbitMQ.

    When specifying properties for a message, pass in a dict of key value items
    that match the AMQP Basic.Properties specification with a small caveat.

    Due to an overlap in the AMQP specification and the Python keyword
    :code:`type`, the :code:`type` property is referred to as
    :code:`message_type`.

    The following is a list of the available properties:

    * app_id
    * content_type
    * content_encoding
    * correlation_id
    * delivery_node
    * expiration
    * headers
    * message_id
    * message_type
    * priority
    * reply_to
    * timestamp
    * user_id

    **Automated features**

    When passing in the body value, if it is a dict or list, it will
    automatically be JSON serialized and the content type "application/json"
    will be set on the message properties.

    When publishing a message to RabbitMQ, if the auto_id value is True and no
    message_id value was passed in as a property, a UUID will be generated and
    specified as a property of the message.

    If a timestamp is not specified when passing in properties, the current
    Unix epoch value will be set in the message properties.

    :param channel: The channel object for the message object to act upon
    :type channel: :py:class:`rabbitpy.channel.Channel`
    :param str or dict or list body_value: The message body
    :param dict properties: A dictionary of message properties
    :param bool auto_id: Add a message id if no properties were passed in.

    """
    method = None
    name = 'Message'

    def __init__(self, channel, body_value, properties=None, auto_id=True):
        """Create a new instance of the Message object."""
        super(Message, self).__init__(channel, 'Message')

        if isinstance(body_value, dict) or isinstance(body_value, list):
            body_value = json.dumps(body_value, ensure_ascii=False)
            if properties is None:
                properties = {'content_type': 'application/json'}
            else:
                properties['content_type'] = 'application/json'
        self.body = body_value
        self.properties = properties or self._base_properties
        if auto_id and 'message_id' not in self.properties:
            self._add_auto_message_id()
        if 'timestamp' not in self.properties:
            self._add_timestamp()
        if isinstance(self.properties['timestamp'], time.struct_time):
            self.properties['timestamp'] = self._timestamp_from_struct_time()

    def _add_auto_message_id(self):
        """Set the message_id property to a new UUID."""
        self.properties['message_id'] = str(uuid.uuid4())

    def _add_timestamp(self):
        """Add the timestamp to the properties"""
        self.properties['timestamp'] = datetime.datetime.now()

    def _timestamp_from_struct_time(self):
        """Return a datetime.datetime object from struct_time

        :rtype: datetime.datetime

        """
        return datetime.datetime(*self.properties['timestamp'][:6])

    @property
    def delivery_tag(self):
        """Return the delivery tag for a message that was delivered or gotten
        from RabbitMQ.

        :rtype: int or None

        """
        return self.method.delivery_tag if self.method else None

    @property
    def _base_properties(self):
        """Return a base set of properties if no properties were passed into
        the constructor.

        :rtype: dict

        """
        return {"message_id": str(uuid.uuid4()),
                "timestamp": datetime.datetime.now()}

    @property
    def _properties(self):
        """Return a new Basic.Properties object representing the message
        properties.

        :rtype: pamqp.specification.Basic.Properties

        """
        invalid_keys = [key for key in self.properties
                        if key not in
                        specification.Basic.Properties.attributes]
        self._prune_invalid_properties(invalid_keys)
        self._coerce_properties()
        return specification.Basic.Properties(**self.properties)

    def _coerce_properties(self):
        """Force properties to be set to the correct data type"""
        for key in self.properties:
            _type = getattr(specification.Basic.Properties, key)
            if DEBUG:
                LOGGER.debug('Type: %s, %s', _type, type(self.properties[key]))
            if _type == 'shortstr' and not isinstance(self.properties[key],
                                                      basestring):
                LOGGER.warning('Coercing property %s to bytes', key)
                self.properties[key] = bytes(self.properties[key])
            elif _type == 'octet' and not isinstance(self.properties[key],
                                                     int):
                LOGGER.warning('Coercing property %s to int', key)
                self.properties[key] = int(self.properties[key])
            elif _type == 'table' and not isinstance(self.properties[key],
                                                     dict):
                LOGGER.warning('Resetting invalid value for %s to None', key)
                self.properties[key] = {}
            elif _type == 'timestamp' and isinstance(self.properties[key],
                                                     int):
                LOGGER.warning('Coercing property %s to datetime', key)
                self.properties[key] = \
                    datetime.datetime.fromtimestamp(self.properties[key])

    def _prune_invalid_properties(self, invalid_keys):
        """Remove invalid properties from the message properties.

        :param list invalid_keys: A list of invalid property names to remove

        """
        for key in invalid_keys:
            LOGGER.warning('Removing invalid property "%s"', key)
            del self.properties[key]

    def ack(self, all_previous=False):
        """Acknowledge receipt of the message to RabbitMQ. Will raise an
        ActionException if the message was not received from a broker.

        :raises: ActionException

        """
        if not self.method:
            raise exceptions.ActionException('Can not ack non-received '
                                             'message')
        basic_ack = specification.Basic.Ack(self.method.delivery_tag,
                                            multiple=all_previous)
        self.channel._write_frame(basic_ack)

    def json(self):
        """Deserialize the message body if it is JSON, returning the value.

        :rtype: any

        """
        return json.loads(self.body)

    def nack(self, requeue=False, all_previous=False):
        """Negatively acknowledge receipt of the message to RabbitMQ. Will
        raise an ActionException if the message was not received from a broker.

        :param bool requeue: Requeue the message
        :param bool all_previous: Nack all previous unacked messages up to and
                                  including this one
        :raises: ActionException

        """
        if not self.method:
            raise exceptions.ActionException('Can not nack non-received '
                                             'message')
        basic_nack = specification.Basic.Nack(self.method.delivery_tag,
                                              requeue=requeue,
                                              multiple=all_previous)
        self.channel._write_frame(basic_nack)

    def pprint(self, properties=False):
        """Print a formatted representation of the message.

        :param bool properties: Include properties in the representation

        """
        print('Exchange: %s\n' % self.method.exchange)
        print('Routing Key: %s\n' % self.method.routing_key)
        if properties:
            print('Properties:\n')
            pprint.pprint(self.properties)
            print('\nBody:\n')
        pprint.pprint(self.body)

    def publish(self, exchange, routing_key='', mandatory=False):
        """Publish the message to the exchange with the specified routing
        key.

        :param str | rabbitpy.base.AMQPClass exchange: The exchange to bind to
        :param str routing_key: The routing key to use
        :param bool mandatory: Requires the message is published
        :return: bool | None
        :raises: rabbitpy.exceptions.MessageReturnedException

        """
        if isinstance(exchange, base.AMQPClass):
            exchange = exchange.name
        method_frame = specification.Basic.Publish(exchange=exchange,
                                                   routing_key=
                                                   routing_key or '',
                                                   mandatory=mandatory)
        self.channel._write_frame(method_frame)
        header_frame = header.ContentHeader(body_size=len(self.body),
                                            properties=self._properties)
        self.channel._write_frame(header_frame)
        pieces = int(math.ceil(len(self.body) /
                               float(self.channel.maximum_frame_size)))

        # Send the message
        for offset in range(0, pieces):
            start = self.channel.maximum_frame_size * offset
            end = start + self.channel.maximum_frame_size
            if end > len(self.body):
                end = len(self.body)
            self.channel._write_frame(body.ContentBody(self.body[start:end]))

        # If publisher confirmations are enabled, wait for the response
        if self.channel.publisher_confirms:
            response = self.channel._wait_for_confirmation()
            if isinstance(response, specification.Basic.Ack):
                return True
            elif isinstance(response, specification.Basic.Nack):
                return False
            else:
                raise exceptions.UnexpectedResponseError(response)

    def reject(self, requeue=False):
        """Reject receipt of the message to RabbitMQ. Will raise
        an ActionException if the message was not received from a broker.

        :param bool requeue: Requeue the message
        :raises: ActionException

        """
        if not self.method:
            raise exceptions.ActionException('Can not reject non-received '
                                             'message')
        basic_reject = specification.Basic.Reject(self.method.delivery_tag,
                                                  requeue=requeue)
        self.channel._write_frame(basic_reject)
<|endoftext|>


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

"""Parse the history file for just about anything."""
import glob
import operator
import os
import os.path
import re
import sys
import time
import threading

def sourceToTitle(source):
  """Convert a source tuple to a title."""
  (short_name, full_name, num_hosts) = source
  return '%s (%s)' % (full_name, num_hosts)

class ParserThread(threading.Thread):
  def __init__(self, record_type):
    threading.Thread.__init__(self)
    self.type = record_type

  def run(self):
    hp = HistoryParser()
    self.hosts = hp.ParseByType(self.type)
    return self.hosts

class HistoryParser(object):
  """Parse the history file from files and web browsers and such."""

  MAX_NON_UNIQUE_RECORD_COUNT = 500000
  MIN_FILE_SIZE = 128
  MIN_RECOMMENDED_RECORD_COUNT = 2
  INTERNAL_RE = re.compile('\.prod|\.corp|\.bor|internal|dmz')
  IP_RE = re.compile('^[\d\.]+$')
  TYPES = {}

  def __init__(self):
    self.TYPES = {
        'chrome': ('Google Chrome', self.GoogleChromeHistoryPath),
        'chromium': ('Chromium', self.ChromiumHistoryPath),
        'epiphany': ('Epiphany', self.EpiphanyHistoryPath),
        'opera': ('Opera', self.OperaHistoryPath),
        'safari': ('Apple Safari', self.SafariHistoryPath),
        'firefox': ('Mozilla Firefox', self.FirefoxHistoryPath),
        'internet_explorer': ('Microsoft Internet Explorer', self.InternetExplorerHistoryPath),
        'konqueror': ('Konqueror', self.KonquerorHistoryPath),
        'seamonkey': ('Mozilla Seamonkey', self.SeamonkeyHistoryPath),
        'squid': ('Squid Web Proxy', self.SquidLogPath),
    }

  def GetTypes(self):
    """Return a tuple of type names with a description."""
    return dict([(x, self.TYPES[x][0]) for x in self.TYPES])

  def GetParsedSource(self, type):
    return self.imported_sources[type]

  def GetAvailableHistorySources(self):
    """Seek out and create a list of valid data sources.

    Returns:
      sources: A list of tuples (type, description, record count)
    """
    source_desc = self.GetTypes()
    self.imported_sources = self.ParseAllTypes()
    sources = []
    for src_type in self.imported_sources:
      sources.append(
        (src_type, source_desc[src_type], len(self.imported_sources[src_type]))
      )
    sources.sort(key=operator.itemgetter(2), reverse=True)
    # TODO(tstromberg): Don't hardcode input types.
    sources.append((None, 'Alexa Top Global Domains', 10000))
    return sources

  def GetTypeMethod(self, type):
    return self.TYPES[type][1]

  def Parse(self, path_or_type):
    if path_or_type == 'auto':
      # TODO(tstromberg): complete this method
      all = self.ParseAllTypes()

    if path_or_type.lower() in self.TYPES:
      return self.ParseByType(path_or_type.lower())
    else:
      return self.ParseByFilename(path_or_type)

  def ParseByType(self, source, complain=True):
    """Given a type, parse the newest file and return a list of hosts."""

    (paths, tried) = self.FindGlobPaths(self.GetTypeMethod(source)())
    if not paths:
      if complain:
        print "- %s: no matches in %s" % (source, tried)
      return False
    newest = sorted(paths, key=os.path.getmtime)[-1]
    return self.ParseByFilename(newest)

  def ParseAllTypes(self):
    """For each type we know of, attempt to find and parse each of them.

    Returns:
      dict of type: list of hosts
    """
    start_time = time.time()
    results = {}
    records = 0
    threads = []

    for type in self.GetTypes():
      thread = ParserThread(type)
      thread.start()
      threads.append(thread)

    for thread in threads:
      thread.join()
      if thread.hosts:
        records += len(thread.hosts)

      if thread.hosts and len(thread.hosts) >= self.MIN_RECOMMENDED_RECORD_COUNT:
        results[thread.type] = thread.hosts
      elif thread.hosts == False:
        pass
      elif not thread.hosts:
        print '- No records found in %s' % (thread.type)
      else:
        print '- Ignoring %s (only %s records)' % (thread.type, len(thread.hosts))

    print '- Read %s records from %s in %ss' % (records, ', '.join(results.keys()), time.time() - start_time)
    return results

  def ParseByFilename(self, filename):
    # Only matches http://host.domain type entries (needs at least one subdom)
    parse_re = re.compile('\w+://([\-\w]+\.[\-\w\.]+)')
    print '* Reading %s' % filename
    # binary mode is necessary for running under Windows
    return parse_re.findall(open(filename, 'rb').read())

  def GenerateTestDataFromInput(self, path_or_type):
    hosts = self.Parse(path_or_type)
    if len(hosts) > self.MAX_NON_UNIQUE_RECORD_COUNT:
      sorted_unique = True
    else:
      sorted_unique = False
    return self.GenerateTestData(hosts, sorted_unique=sorted_unique)

  def GenerateTestData(self, hosts, sorted_unique=False):
    """Given a set of hosts, generate test data.

    Args:
      hosts: A list of hostnames
      sorted_unique: Return a sorted unique list of tests. Useful for large
                     data sets.

    Returns:
      A list of strings representing DNS requests to test with.
    """
    history = []
    hits = {}
    last_host = None

    for host in hosts:
      if not host.endswith('.'):
        host = host + '.'

      if self.INTERNAL_RE.search(host):
        continue

      if self.IP_RE.match(host):
        continue

      if host != last_host:
        if sorted_unique:
          hits[host] = hits.get(host, 0) + 1
        else:
          history.append('A %s' % host)
        last_host = host

    if sorted_unique:
      for (hit, count) in sorted(hits.items(), key=operator.itemgetter(1),
                                 reverse=True):
        history.append('A %s # %s hits' % (hit, count))
    return history

  def FindGlobPaths(self, paths):
    """Given a list of glob paths, return a list of matches containing data.

    Returns:
      A tuple with (file path (str), list of paths checked)
    """
    tried = []
    found = []
    for path_elements in paths:
      path = os.path.join(*path_elements)
      tried.append(path)
      for filename in glob.glob(path):
        if os.path.getsize(filename) > self.MIN_FILE_SIZE:
          found.append(filename)
        else:
          print "* %s exists, but is only %s bytes" % (filename, os.path.getsize(filename))

    return (found, tried)

  def GoogleChromeHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Application Support', 'Google',
         'Chrome', 'Default', 'History'),
        (os.getenv('HOME', ''), '.config', 'google-chrome', 'Default',
         'History'),
        (os.getenv('APPDATA', ''), 'Google', 'Chrome', 'User Data', 'Default',
         'History'),
        (os.getenv('APPDATA', ''), 'Local', 'Google', 'Chrome', 'User Data',
         'Default', 'History'),
        (os.getenv('USERPROFILE', ''), 'Local Settings', 'Application Data',
         'Google', 'Chrome', 'User Data', 'Default', 'History'),
    )
    return paths

  def ChromiumHistoryPath(self):
    """It's like Chrome, but with the branding stripped out."""

    # TODO(tstromberg): Find a terser way to do this.
    paths = []
    for path in self.GoogleChromeHistoryPath():
      new_path = list(path)
      if 'Google' in new_path:
        new_path.remove('Google')
      for (index, part) in enumerate(new_path):
        if part == 'Chrome':
          new_path[index] = 'Chromium'
        elif part == 'chrome' or part == 'google-chrome':
          new_path[index] = 'chromium'
      paths.append(new_path)

    return paths

  def OperaHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Preferences', 'Opera Preferences',
         'global_history.dat'),
        (os.getenv('HOME', ''), 'Library', 'Preferences', 'Opera Preferences 10',
         'global_history.dat'),
        (os.getenv('APPDATA', ''), 'Opera', 'Opera', 'global_history.dat'),
        (os.getenv('APPDATA', ''), 'Local', 'Opera', 'Opera', 'global_history.dat'),
        (os.getenv('HOME', ''), '.opera', 'global_history.dat'),
    )
    return paths

  def SafariHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Safari', 'History.plist'),
        (os.getenv('APPDATA', ''), 'Apple Computer', 'Safari',
         'History.plist')
    )
    return paths

  def FirefoxHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Application Support', 'Firefox',
         'Profiles', '*', 'places.sqlite'),
        (os.getenv('HOME', ''), '.mozilla', 'firefox', '*', 'places.sqlite'),
        (os.getenv('APPDATA', ''), 'Mozilla', 'Firefox', 'Profiles', '*',
         'places.sqlite')
    )
    return paths

  def SeamonkeyHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Application Support', 'Seamonkey',
         'Profiles', '*', 'history.dat'),
        (os.getenv('HOME', ''), '.mozilla', 'seamonkey', '*', 'history.dat'),
        (os.getenv('HOME', ''), '.mozilla', 'default', '*', 'history.dat'),
        (os.getenv('APPDATA', ''), 'Mozilla', 'Seamonkey', 'Profiles', '*',
         'history.dat')
    )
    return paths

  def InternetExplorerHistoryPath(self):
    paths = (
        # XP
        (os.getenv('USERPROFILE', ''), 'Local Settings', 'History',
         'History.IE5', 'index.dat'),
        # ?
        (os.getenv('APPDATA', ''), 'Microsoft', 'Windows', 'History',
         'History.IE5', 'index.dat'),
    )
    return paths

  def EpiphanyHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), '.gnome2', 'epiphany', 'ephy-history.xml'),
    )
    return paths

  def KonquerorHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), '.kde4', 'share',  'apps', 'konqueror', 'konq_history'),
        (os.getenv('HOME', ''), '.kde', 'share',  'apps', 'konqueror', 'konq_history'),
    )
    return paths

  def SquidLogPath(self):
    paths = (
        ('/usr/local/squid/logs/access.log',),
        ('/var/log/squid/access_log',)
    )
    return paths

if __name__ == '__main__':
  parser = HistoryParser()
  types_str = ', '.join(parser.TYPES.keys())
  if len(sys.argv) < 2:
    print 'You must provide a filename or history file type (%s)' % types_str
    sys.exit(1)

  records = parser.Parse(sys.argv[1])
  for record in records:
    print record

<commit_msg>Add another debug message to history_parser to show counts per source<commit_after>#!/usr/bin/env python
# Copyright 2009 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Parse the history file for just about anything."""
import glob
import operator
import os
import os.path
import re
import sys
import time
import threading

def sourceToTitle(source):
  """Convert a source tuple to a title."""
  (short_name, full_name, num_hosts) = source
  return '%s (%s)' % (full_name, num_hosts)

class ParserThread(threading.Thread):
  def __init__(self, record_type):
    threading.Thread.__init__(self)
    self.type = record_type

  def run(self):
    hp = HistoryParser()
    self.hosts = hp.ParseByType(self.type)
    return self.hosts

class HistoryParser(object):
  """Parse the history file from files and web browsers and such."""

  MAX_NON_UNIQUE_RECORD_COUNT = 500000
  MIN_FILE_SIZE = 128
  MIN_RECOMMENDED_RECORD_COUNT = 2
  INTERNAL_RE = re.compile('\.prod|\.corp|\.bor|internal|dmz')
  IP_RE = re.compile('^[\d\.]+$')
  TYPES = {}

  def __init__(self):
    self.TYPES = {
        'chrome': ('Google Chrome', self.GoogleChromeHistoryPath),
        'chromium': ('Chromium', self.ChromiumHistoryPath),
        'epiphany': ('Epiphany', self.EpiphanyHistoryPath),
        'opera': ('Opera', self.OperaHistoryPath),
        'safari': ('Apple Safari', self.SafariHistoryPath),
        'firefox': ('Mozilla Firefox', self.FirefoxHistoryPath),
        'internet_explorer': ('Microsoft Internet Explorer', self.InternetExplorerHistoryPath),
        'konqueror': ('Konqueror', self.KonquerorHistoryPath),
        'seamonkey': ('Mozilla Seamonkey', self.SeamonkeyHistoryPath),
        'squid': ('Squid Web Proxy', self.SquidLogPath),
    }

  def GetTypes(self):
    """Return a tuple of type names with a description."""
    return dict([(x, self.TYPES[x][0]) for x in self.TYPES])

  def GetParsedSource(self, type):
    return self.imported_sources[type]

  def GetAvailableHistorySources(self):
    """Seek out and create a list of valid data sources.

    Returns:
      sources: A list of tuples (type, description, record count)
    """
    source_desc = self.GetTypes()
    self.imported_sources = self.ParseAllTypes()
    sources = []
    for src_type in self.imported_sources:
      sources.append(
        (src_type, source_desc[src_type], len(self.imported_sources[src_type]))
      )
    sources.sort(key=operator.itemgetter(2), reverse=True)
    # TODO(tstromberg): Don't hardcode input types.
    sources.append((None, 'Alexa Top Global Domains', 10000))
    return sources

  def GetTypeMethod(self, type):
    return self.TYPES[type][1]

  def Parse(self, path_or_type):
    if path_or_type == 'auto':
      # TODO(tstromberg): complete this method
      all = self.ParseAllTypes()

    if path_or_type.lower() in self.TYPES:
      return self.ParseByType(path_or_type.lower())
    else:
      return self.ParseByFilename(path_or_type)

  def ParseByType(self, source, complain=True):
    """Given a type, parse the newest file and return a list of hosts."""

    (paths, tried) = self.FindGlobPaths(self.GetTypeMethod(source)())
    if not paths:
      if complain:
        print "- %s: no matches in %s" % (source, tried)
      return False
    newest = sorted(paths, key=os.path.getmtime)[-1]
    return self.ParseByFilename(newest)

  def ParseAllTypes(self):
    """For each type we know of, attempt to find and parse each of them.

    Returns:
      dict of type: list of hosts
    """
    start_time = time.time()
    results = {}
    records = 0
    threads = []

    for type in self.GetTypes():
      thread = ParserThread(type)
      thread.start()
      threads.append(thread)

    for thread in threads:
      thread.join()
      if thread.hosts:
        records += len(thread.hosts)

      if thread.hosts and len(thread.hosts) >= self.MIN_RECOMMENDED_RECORD_COUNT:
        results[thread.type] = thread.hosts
        print "- Found %s records from %s" % (len(thread.hosts), thread.type)
      elif thread.hosts == False:
        pass
      elif not thread.hosts:
        print '- No records found in %s' % (thread.type)
      else:
        print '- Ignoring %s (only %s records)' % (thread.type, len(thread.hosts))

    print '- Read %s records from %s in %ss' % (records, ', '.join(results.keys()), time.time() - start_time)
    return results

  def ParseByFilename(self, filename):
    # Only matches http://host.domain type entries (needs at least one subdom)
    parse_re = re.compile('\w+://([\-\w]+\.[\-\w\.]+)')
    print '* Reading %s' % filename
    # binary mode is necessary for running under Windows
    return parse_re.findall(open(filename, 'rb').read())

  def GenerateTestDataFromInput(self, path_or_type):
    hosts = self.Parse(path_or_type)
    if len(hosts) > self.MAX_NON_UNIQUE_RECORD_COUNT:
      sorted_unique = True
    else:
      sorted_unique = False
    return self.GenerateTestData(hosts, sorted_unique=sorted_unique)

  def GenerateTestData(self, hosts, sorted_unique=False):
    """Given a set of hosts, generate test data.

    Args:
      hosts: A list of hostnames
      sorted_unique: Return a sorted unique list of tests. Useful for large
                     data sets.

    Returns:
      A list of strings representing DNS requests to test with.
    """
    history = []
    hits = {}
    last_host = None

    for host in hosts:
      if not host.endswith('.'):
        host = host + '.'

      if self.INTERNAL_RE.search(host):
        continue

      if self.IP_RE.match(host):
        continue

      if host != last_host:
        if sorted_unique:
          hits[host] = hits.get(host, 0) + 1
        else:
          history.append('A %s' % host)
        last_host = host

    if sorted_unique:
      for (hit, count) in sorted(hits.items(), key=operator.itemgetter(1),
                                 reverse=True):
        history.append('A %s # %s hits' % (hit, count))
    return history

  def FindGlobPaths(self, paths):
    """Given a list of glob paths, return a list of matches containing data.

    Returns:
      A tuple with (file path (str), list of paths checked)
    """
    tried = []
    found = []
    for path_elements in paths:
      path = os.path.join(*path_elements)
      tried.append(path)
      for filename in glob.glob(path):
        if os.path.getsize(filename) > self.MIN_FILE_SIZE:
          found.append(filename)
        else:
          print "* %s exists, but is only %s bytes" % (filename, os.path.getsize(filename))

    return (found, tried)

  def GoogleChromeHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Application Support', 'Google',
         'Chrome', 'Default', 'History'),
        (os.getenv('HOME', ''), '.config', 'google-chrome', 'Default',
         'History'),
        (os.getenv('APPDATA', ''), 'Google', 'Chrome', 'User Data', 'Default',
         'History'),
        (os.getenv('APPDATA', ''), 'Local', 'Google', 'Chrome', 'User Data',
         'Default', 'History'),
        (os.getenv('USERPROFILE', ''), 'Local Settings', 'Application Data',
         'Google', 'Chrome', 'User Data', 'Default', 'History'),
    )
    return paths

  def ChromiumHistoryPath(self):
    """It's like Chrome, but with the branding stripped out."""

    # TODO(tstromberg): Find a terser way to do this.
    paths = []
    for path in self.GoogleChromeHistoryPath():
      new_path = list(path)
      if 'Google' in new_path:
        new_path.remove('Google')
      for (index, part) in enumerate(new_path):
        if part == 'Chrome':
          new_path[index] = 'Chromium'
        elif part == 'chrome' or part == 'google-chrome':
          new_path[index] = 'chromium'
      paths.append(new_path)

    return paths

  def OperaHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Preferences', 'Opera Preferences',
         'global_history.dat'),
        (os.getenv('HOME', ''), 'Library', 'Preferences', 'Opera Preferences 10',
         'global_history.dat'),
        (os.getenv('APPDATA', ''), 'Opera', 'Opera', 'global_history.dat'),
        (os.getenv('APPDATA', ''), 'Local', 'Opera', 'Opera', 'global_history.dat'),
        (os.getenv('HOME', ''), '.opera', 'global_history.dat'),
    )
    return paths

  def SafariHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Safari', 'History.plist'),
        (os.getenv('APPDATA', ''), 'Apple Computer', 'Safari',
         'History.plist')
    )
    return paths

  def FirefoxHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Application Support', 'Firefox',
         'Profiles', '*', 'places.sqlite'),
        (os.getenv('HOME', ''), '.mozilla', 'firefox', '*', 'places.sqlite'),
        (os.getenv('APPDATA', ''), 'Mozilla', 'Firefox', 'Profiles', '*',
         'places.sqlite')
    )
    return paths

  def SeamonkeyHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), 'Library', 'Application Support', 'Seamonkey',
         'Profiles', '*', 'history.dat'),
        (os.getenv('HOME', ''), '.mozilla', 'seamonkey', '*', 'history.dat'),
        (os.getenv('HOME', ''), '.mozilla', 'default', '*', 'history.dat'),
        (os.getenv('APPDATA', ''), 'Mozilla', 'Seamonkey', 'Profiles', '*',
         'history.dat')
    )
    return paths

  def InternetExplorerHistoryPath(self):
    paths = (
        # XP
        (os.getenv('USERPROFILE', ''), 'Local Settings', 'History',
         'History.IE5', 'index.dat'),
        # ?
        (os.getenv('APPDATA', ''), 'Microsoft', 'Windows', 'History',
         'History.IE5', 'index.dat'),
    )
    return paths

  def EpiphanyHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), '.gnome2', 'epiphany', 'ephy-history.xml'),
    )
    return paths

  def KonquerorHistoryPath(self):
    paths = (
        (os.getenv('HOME', ''), '.kde4', 'share',  'apps', 'konqueror', 'konq_history'),
        (os.getenv('HOME', ''), '.kde', 'share',  'apps', 'konqueror', 'konq_history'),
    )
    return paths

  def SquidLogPath(self):
    paths = (
        ('/usr/local/squid/logs/access.log',),
        ('/var/log/squid/access_log',)
    )
    return paths

if __name__ == '__main__':
  parser = HistoryParser()
  types_str = ', '.join(parser.TYPES.keys())
  if len(sys.argv) < 2:
    print 'You must provide a filename or history file type (%s)' % types_str
    sys.exit(1)

  records = parser.Parse(sys.argv[1])
  for record in records:
    print record

<|endoftext|>


<commit_before>from .null_model import NullModel
from sklearn import linear_model as lm

def linear_model(context, node):
    partition = node.data
    if partition.y.size < 2:
        return NullModel()
    model = lm.LinearRegression()
    model.fit(partition.x, partition.y)
    return model

def model_of(klass, **kwargs):
    def _create(context, node):
        partition = node.data
        if partition.y.size < 2:
            return
        model = klass(**kwargs)
        model.fit(partition.x, partition.y)
        return model

    return _create

def dim_models(context, node):
    partition = node.data
    if partition.y.size < 2:
        return [NullModel() for d in partition.x.columns]

    models = [lm.LinearRegression().fit(partition.x[[d]], partition.y) for d in partition.x.columns]
    return models
<commit_msg>add ridge model<commit_after>from .null_model import NullModel
from sklearn import linear_model as lm

def linear_model(context, node):
    partition = node.data
    if partition.y.size < 2:
        return NullModel()
    model = lm.LinearRegression()
    model.fit(partition.x, partition.y)
    return model

def ridge_model(context, node):
    partition = node.data
    if partition.y.size < 2:
        return NullModel()
    model = lm.Ridge(alpha=1.0)
    model.fit(partition.x, partition.y)
    return model

def model_of(klass, **kwargs):
    def _create(context, node):
        partition = node.data
        if partition.y.size < 2:
            return
        model = klass(**kwargs)
        model.fit(partition.x, partition.y)
        return model

    return _create

def dim_models(context, node):
    partition = node.data
    if partition.y.size < 2:
        return [NullModel() for d in partition.x.columns]

    models = [lm.LinearRegression().fit(partition.x[[d]], partition.y) for d in partition.x.columns]
    return models
<|endoftext|>


<commit_before># These three constants are used to define the mode that a worker is running in.
# Right now, this is only used for determining how to print information about
# task failures.
SCRIPT_MODE = 0
WORKER_MODE = 1
SHELL_MODE = 2
PYTHON_MODE = 3

import config
import libraylib as lib
import serialization
from worker import scheduler_info, visualize_computation_graph, task_info, register_module, connect, disconnect, get, put, remote, kill_workers, restart_workers_local
from libraylib import ObjRef
import internal
<commit_msg>Kill child processes when parent dies, and die on fatal errors (#265)<commit_after># These three constants are used to define the mode that a worker is running in.
# Right now, this is only used for determining how to print information about
# task failures.
SCRIPT_MODE = 0
WORKER_MODE = 1
SHELL_MODE = 2
PYTHON_MODE = 3

import ctypes
# Windows only
if hasattr(ctypes, "windll"):
  # Makes sure that all child processes die when we die
  # Also makes sure that fatal crashes result in process termination rather than an error dialog (the latter is annoying since we have a lot of processes)
  # This is done by associating all child processes with a "job" object that imposes this behavior.
  (lambda kernel32: (lambda job: (lambda n: kernel32.SetInformationJobObject(job, 9, "\0" * 17 + chr(0x8 | 0x4 | 0x20) + "\0" * (n - 18), n))(0x90 if ctypes.sizeof(ctypes.c_void_p) > ctypes.sizeof(ctypes.c_int) else 0x70) and kernel32.AssignProcessToJobObject(job, ctypes.c_void_p(kernel32.GetCurrentProcess())))(ctypes.c_void_p(kernel32.CreateJobObjectW(None, None))) if kernel32 is not None else None)(ctypes.windll.kernel32)

import config
import libraylib as lib
import serialization
from worker import scheduler_info, visualize_computation_graph, task_info, register_module, connect, disconnect, get, put, remote, kill_workers, restart_workers_local
from libraylib import ObjRef
import internal
<|endoftext|>