<commit_before># Copyright 2019 The Johns Hopkins University Applied Physics Laboratory
#
# 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.

from lib.ssh import SSHConnection, SSHTarget
from lib import utils
from lib import constants as const
from lib import zip

import botocore
import configparser
import os
import tempfile
import subprocess
import shlex
import shutil

# Location of settings files for ndingest.
NDINGEST_SETTINGS_FOLDER = const.repo_path('salt_stack', 'salt', 'ndingest', 'files', 'ndingest.git', 'settings')

# Template used for ndingest settings.ini generation.
NDINGEST_SETTINGS_TEMPLATE = NDINGEST_SETTINGS_FOLDER + '/settings.ini.apl'

def update_lambda_code(bosslet_config):
    """Update all lambdas that use the multilambda zip file.

    Args:
        bosslet_config: Bosslet configuration object
    """
    names = bosslet_config.names
    uses_multilambda = [
        names.multi_lambda.lambda_,
        names.downsample_volume.lambda_,
        names.delete_tile_objs.lambda_,
        names.delete_tile_index_entry.lambda_,
        names.index_s3_writer.lambda_,
        names.index_fanout_id_writer.lambda_,
        names.index_write_id.lambda_,
        names.index_write_failed.lambda_,
        names.index_find_cuboids.lambda_,
        names.index_split_cuboids.lambda_,
        names.index_fanout_enqueue_cuboid_keys.lambda_,
        names.index_batch_enqueue_cuboids.lambda_,
        names.index_fanout_dequeue_cuboid_keys.lambda_,
        names.index_dequeue_cuboid_keys.lambda_,
        names.index_get_num_cuboid_keys_msgs.lambda_,
        names.index_check_for_throttling.lambda_,
        names.index_invoke_index_supervisor.lambda_,
        names.index_load_ids_from_s3.lambda_,
        names.start_sfn.lambda_,
        names.copy_cuboid_lambda.lambda_,
        names.cuboid_import_lambda.lambda_,
        names.volumetric_ingest_queue_upload_lambda.lambda_,
        names.tile_uploaded.lambda_,
        names.tile_ingest.lambda_,
        names.delete_tile_index_entry.lambda_,
    ]
    client = bosslet_config.session.client('lambda')
    for lambda_name in uses_multilambda:
        try:
            resp = client.update_function_code(
                FunctionName=lambda_name,
                S3Bucket=bosslet_config.LAMBDA_BUCKET,
                S3Key=names.multi_lambda.zip,
                Publish=True)
            print(resp)
        except botocore.exceptions.ClientError as ex:
            print('Error updating {}: {}'.format(lambda_name, ex))

def load_lambdas_on_s3(bosslet_config):
    """Zip up spdb, bossutils, lambda and lambda_utils.  Upload to S3.

    Uses the lambda build server (an Amazon Linux AMI) to compile C code and
    prepare the virtualenv that's ultimately contained in the zip file placed
    in S3.

    Args:
        session (Session): boto3.Session
        domain (str): The VPC's domain name such as integration.boss.
        bucket (str): Name of bucket that contains the lambda zip file.
    """
    domain = bosslet_config.INTERNAL_DOMAIN
    tempname = tempfile.NamedTemporaryFile(delete=True)
    zipname = tempname.name + '.zip'
    tempname.close()
    print('Using temp zip file: ' + zipname)

    cwd = os.getcwd()
    os.chdir(const.repo_path("salt_stack", "salt", "spdb", "files"))
    zip.write_to_zip('spdb.git', zipname, False)
    os.chdir(cwd)

    os.chdir(const.repo_path("salt_stack", "salt", "boss-tools", "files", "boss-tools.git"))
    zip.write_to_zip('bossutils', zipname)
    zip.write_to_zip('cloudwatchwrapper', zipname)
    zip.write_to_zip('lambda', zipname)
    zip.write_to_zip('lambdautils', zipname)
    os.chdir(cwd)

    with open(NDINGEST_SETTINGS_TEMPLATE, 'r') as tmpl:
        # Generate settings.ini file for ndingest.
        create_ndingest_settings(bosslet_config, tmpl)

    os.chdir(const.repo_path("salt_stack", "salt", "ndingest", "files"))
    zip.write_to_zip('ndingest.git', zipname)
    os.chdir(cwd)

    os.chdir(const.repo_path("lib"))
    zip.write_to_zip('heaviside.git', zipname)

    # Let lambdas look up names by creating a bossnames module.
    zip.write_to_zip('names.py', zipname, arcname='bossnames/names.py')
    zip.write_to_zip('bucket_object_tags.py', zipname, arcname='bossnames/bucket_object_tags.py')
    zip.write_to_zip('__init__.py', zipname, arcname='bossnames/__init__.py')
    os.chdir(cwd)

    #copy the zip file to lambda_build_server and run makedomainenv
    lambda_bucket = bosslet_config.LAMBDA_BUCKET

    build_cmd = 'source /etc/profile && source ~/.bash_profile && /home/ec2-user/makedomainenv {} {}'.format(domain, lambda_bucket)
    target_file = "sitezips/{}.zip".format(domain)

    lambda_build_server = bosslet_config.LAMBDA_SERVER
    if lambda_build_server is None:
        print("Copying local modules to localhost")
        full_target_file = '/home/ec2-user/' + target_file
        shutil.copy(zipname, full_target_file)
        os.remove(zipname)

        try:
            print("calling makedomainenv on localhost")

            output = subprocess.check_output(build_cmd,
                                             shell=True,
                                             executable='/bin/bash',
                                             stderr=subprocess.STDOUT)
            print(output.decode('utf-8'))
        except subprocess.CalledProcessError as ex:
            print("makedomainenv return code: {}".format(ex.returncode))
            print(ex.output.decode('utf-8'))
            # DP NOTE: currently eating the error, as there is no checking of error if there is a build server
        finally:
            os.remove(full_target_file)
    else:
        lambda_build_server_key = bosslet_config.LAMBDA_SERVER_KEY
        lambda_build_server_key = utils.keypair_to_file(lambda_build_server_key)
        ssh_target = SSHTarget(lambda_build_server_key, lambda_build_server, 22, 'ec2-user')
        bastions = [bosslet_config.outbound_bastion] if bosslet_config.outbound_bastion else []
        ssh = SSHConnection(ssh_target, bastions)

        print("Copying local modules to lambda-build-server")
        ret = ssh.scp(zipname, target_file, upload=True)
        print("scp return code: " + str(ret))

        os.remove(zipname)

        # This section will run makedomainenv on lambda-build-server
        print("calling makedomainenv on lambda-build-server")
        ret = ssh.cmd(build_cmd)
        print("ssh return code: " + str(ret))

def create_ndingest_settings(bosslet_config, fp):
    """Create the settings.ini file for ndingest.

    The file is placed in ndingest's settings folder.

    Args:
        domain (str): The VPC's domain name such as integration.boss.
        fp (file-like object): File like object to read settings.ini template from.
    """
    names = bosslet_config.names
    parser = configparser.ConfigParser()
    parser.read_file(fp)

    parser['boss']['domain'] = bosslet_config.INTERNAL_DOMAIN

    parser['aws']['region'] = bosslet_config.REGION
    parser['aws']['tile_bucket'] = names.tile_bucket.s3
    parser['aws']['cuboid_bucket'] = names.cuboid_bucket.s3
    parser['aws']['tile_index_table'] = names.tile_index.ddb
    parser['aws']['cuboid_index_table'] = names.s3_index.ddb
    parser['aws']['max_task_id_suffix'] = str(const.MAX_TASK_ID_SUFFIX)

    # parser['spdb']['SUPER_CUBOID_SIZE'] = CUBOIDSIZE[0]
    # ToDo: find way to always get cuboid size from spdb.
    parser['spdb']['SUPER_CUBOID_SIZE'] = '512, 512, 16'

    with open(NDINGEST_SETTINGS_FOLDER + '/settings.ini', 'w') as out:
        parser.write(out)

def freshen_lambda(bosslet_config, lambda_name):
    """
    Tell a lambda to reload its code from S3.

    Useful when developing and small changes need to be made to a lambda function,
    but a full rebuild of the entire zip file isn't required.
    """
    zip_name = bosslet_config.names.multi_lambda.zip
    full_name = bosslet_config.names[lambda_name].lambda_
    client = bosslet_config.session.client('lambda')
    resp = client.update_function_code(
        FunctionName=full_name,
        S3Bucket=bosslet_config.LAMBDA_BUCKET,
        S3Key=zip_name,
        Publish=True)
    print(resp)

def download_lambda_zip(bosslet_config, path):
    """
    Download the existing multilambda.domain.zip from the S3 bucket.  Useful when
    developing and small changes need to be made to a lambda function, but a full
    rebuild of the entire zip file isn't required.
    """
    s3 = bosslet_config.session.client('s3')
    zip_name = bosslet_config.names.multi_lambda.zip
    full_path = '{}/{}'.format(path, zip_name)
    resp = s3.get_object(Bucket=bosslet_config.LAMBDA_BUCKET, Key=zip_name)

    bytes = resp['Body'].read()

    with open(full_path , 'wb') as out:
        out.write(bytes)

    print('Saved zip to {}'.format(full_path))

def upload_lambda_zip(bosslet_config, path):
    """
    Upload a  multilambda.domain.zip to the S3 bucket.  Useful when
    developing and small changes need to be made to a lambda function, but a full
    rebuild of the entire zip file isn't required.
    """
    s3 = bosslet_config.session.client('s3')
    with open(path, 'rb') as in_file:
        resp = s3.put_object(Bucket=bosslet_config.LAMBDA_BUCKET,
                             Key=bosslet_config.names.multi_lambda.zip,
                             Body=in_file)
    print(resp)

<commit_msg>Updated local lambda build so that it will us the ec2-user, so that the build scripts work correctly<commit_after># Copyright 2019 The Johns Hopkins University Applied Physics Laboratory
#
# 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.

from lib.ssh import SSHConnection, SSHTarget
from lib import utils
from lib import constants as const
from lib import zip

import botocore
import configparser
import os
import tempfile
import subprocess
import shlex
import shutil
import pwd

# Location of settings files for ndingest.
NDINGEST_SETTINGS_FOLDER = const.repo_path('salt_stack', 'salt', 'ndingest', 'files', 'ndingest.git', 'settings')

# Template used for ndingest settings.ini generation.
NDINGEST_SETTINGS_TEMPLATE = NDINGEST_SETTINGS_FOLDER + '/settings.ini.apl'

def update_lambda_code(bosslet_config):
    """Update all lambdas that use the multilambda zip file.

    Args:
        bosslet_config: Bosslet configuration object
    """
    names = bosslet_config.names
    uses_multilambda = [
        names.multi_lambda.lambda_,
        names.downsample_volume.lambda_,
        names.delete_tile_objs.lambda_,
        names.delete_tile_index_entry.lambda_,
        names.index_s3_writer.lambda_,
        names.index_fanout_id_writer.lambda_,
        names.index_write_id.lambda_,
        names.index_write_failed.lambda_,
        names.index_find_cuboids.lambda_,
        names.index_split_cuboids.lambda_,
        names.index_fanout_enqueue_cuboid_keys.lambda_,
        names.index_batch_enqueue_cuboids.lambda_,
        names.index_fanout_dequeue_cuboid_keys.lambda_,
        names.index_dequeue_cuboid_keys.lambda_,
        names.index_get_num_cuboid_keys_msgs.lambda_,
        names.index_check_for_throttling.lambda_,
        names.index_invoke_index_supervisor.lambda_,
        names.index_load_ids_from_s3.lambda_,
        names.start_sfn.lambda_,
        names.copy_cuboid_lambda.lambda_,
        names.cuboid_import_lambda.lambda_,
        names.volumetric_ingest_queue_upload_lambda.lambda_,
        names.tile_uploaded.lambda_,
        names.tile_ingest.lambda_,
        names.delete_tile_index_entry.lambda_,
    ]
    client = bosslet_config.session.client('lambda')
    for lambda_name in uses_multilambda:
        try:
            resp = client.update_function_code(
                FunctionName=lambda_name,
                S3Bucket=bosslet_config.LAMBDA_BUCKET,
                S3Key=names.multi_lambda.zip,
                Publish=True)
            print(resp)
        except botocore.exceptions.ClientError as ex:
            print('Error updating {}: {}'.format(lambda_name, ex))

def load_lambdas_on_s3(bosslet_config):
    """Zip up spdb, bossutils, lambda and lambda_utils.  Upload to S3.

    Uses the lambda build server (an Amazon Linux AMI) to compile C code and
    prepare the virtualenv that's ultimately contained in the zip file placed
    in S3.

    Args:
        session (Session): boto3.Session
        domain (str): The VPC's domain name such as integration.boss.
        bucket (str): Name of bucket that contains the lambda zip file.
    """
    domain = bosslet_config.INTERNAL_DOMAIN
    tempname = tempfile.NamedTemporaryFile(delete=True)
    zipname = tempname.name + '.zip'
    tempname.close()
    print('Using temp zip file: ' + zipname)

    cwd = os.getcwd()
    os.chdir(const.repo_path("salt_stack", "salt", "spdb", "files"))
    zip.write_to_zip('spdb.git', zipname, False)
    os.chdir(cwd)

    os.chdir(const.repo_path("salt_stack", "salt", "boss-tools", "files", "boss-tools.git"))
    zip.write_to_zip('bossutils', zipname)
    zip.write_to_zip('cloudwatchwrapper', zipname)
    zip.write_to_zip('lambda', zipname)
    zip.write_to_zip('lambdautils', zipname)
    os.chdir(cwd)

    with open(NDINGEST_SETTINGS_TEMPLATE, 'r') as tmpl:
        # Generate settings.ini file for ndingest.
        create_ndingest_settings(bosslet_config, tmpl)

    os.chdir(const.repo_path("salt_stack", "salt", "ndingest", "files"))
    zip.write_to_zip('ndingest.git', zipname)
    os.chdir(cwd)

    os.chdir(const.repo_path("lib"))
    zip.write_to_zip('heaviside.git', zipname)

    # Let lambdas look up names by creating a bossnames module.
    zip.write_to_zip('names.py', zipname, arcname='bossnames/names.py')
    zip.write_to_zip('bucket_object_tags.py', zipname, arcname='bossnames/bucket_object_tags.py')
    zip.write_to_zip('__init__.py', zipname, arcname='bossnames/__init__.py')
    os.chdir(cwd)

    #copy the zip file to lambda_build_server and run makedomainenv
    lambda_bucket = bosslet_config.LAMBDA_BUCKET

    build_cmd = 'source /etc/profile && source ~/.bash_profile && /home/ec2-user/makedomainenv {} {}'.format(domain, lambda_bucket)
    target_file = "sitezips/{}.zip".format(domain)

    lambda_build_server = bosslet_config.LAMBDA_SERVER
    if lambda_build_server is None:
        print("Copying local modules to localhost")
        full_target_file = '/home/ec2-user/' + target_file
        shutil.copy(zipname, full_target_file)
        os.remove(zipname)

        cur_user = pwd.getpwuid(os.getuid()).pw_name
        if cur_user != 'ec2-user':
            build_cmd = 'su -m ec2-user -c "{}"'.format(build_cmd)

        try:
            print("calling makedomainenv on localhost")

            output = subprocess.check_output(build_cmd,
                                             shell=True,
                                             executable='/bin/bash',
                                             stderr=subprocess.STDOUT)
            print(output.decode('utf-8'))
        except subprocess.CalledProcessError as ex:
            print("makedomainenv return code: {}".format(ex.returncode))
            print(ex.output.decode('utf-8'))
            # DP NOTE: currently eating the error, as there is no checking of error if there is a build server
        finally:
            os.remove(full_target_file)
    else:
        lambda_build_server_key = bosslet_config.LAMBDA_SERVER_KEY
        lambda_build_server_key = utils.keypair_to_file(lambda_build_server_key)
        ssh_target = SSHTarget(lambda_build_server_key, lambda_build_server, 22, 'ec2-user')
        bastions = [bosslet_config.outbound_bastion] if bosslet_config.outbound_bastion else []
        ssh = SSHConnection(ssh_target, bastions)

        print("Copying local modules to lambda-build-server")
        ret = ssh.scp(zipname, target_file, upload=True)
        print("scp return code: " + str(ret))

        os.remove(zipname)

        # This section will run makedomainenv on lambda-build-server
        print("calling makedomainenv on lambda-build-server")
        ret = ssh.cmd(build_cmd)
        print("ssh return code: " + str(ret))

def create_ndingest_settings(bosslet_config, fp):
    """Create the settings.ini file for ndingest.

    The file is placed in ndingest's settings folder.

    Args:
        domain (str): The VPC's domain name such as integration.boss.
        fp (file-like object): File like object to read settings.ini template from.
    """
    names = bosslet_config.names
    parser = configparser.ConfigParser()
    parser.read_file(fp)

    parser['boss']['domain'] = bosslet_config.INTERNAL_DOMAIN

    parser['aws']['region'] = bosslet_config.REGION
    parser['aws']['tile_bucket'] = names.tile_bucket.s3
    parser['aws']['cuboid_bucket'] = names.cuboid_bucket.s3
    parser['aws']['tile_index_table'] = names.tile_index.ddb
    parser['aws']['cuboid_index_table'] = names.s3_index.ddb
    parser['aws']['max_task_id_suffix'] = str(const.MAX_TASK_ID_SUFFIX)

    # parser['spdb']['SUPER_CUBOID_SIZE'] = CUBOIDSIZE[0]
    # ToDo: find way to always get cuboid size from spdb.
    parser['spdb']['SUPER_CUBOID_SIZE'] = '512, 512, 16'

    with open(NDINGEST_SETTINGS_FOLDER + '/settings.ini', 'w') as out:
        parser.write(out)

def freshen_lambda(bosslet_config, lambda_name):
    """
    Tell a lambda to reload its code from S3.

    Useful when developing and small changes need to be made to a lambda function,
    but a full rebuild of the entire zip file isn't required.
    """
    zip_name = bosslet_config.names.multi_lambda.zip
    full_name = bosslet_config.names[lambda_name].lambda_
    client = bosslet_config.session.client('lambda')
    resp = client.update_function_code(
        FunctionName=full_name,
        S3Bucket=bosslet_config.LAMBDA_BUCKET,
        S3Key=zip_name,
        Publish=True)
    print(resp)

def download_lambda_zip(bosslet_config, path):
    """
    Download the existing multilambda.domain.zip from the S3 bucket.  Useful when
    developing and small changes need to be made to a lambda function, but a full
    rebuild of the entire zip file isn't required.
    """
    s3 = bosslet_config.session.client('s3')
    zip_name = bosslet_config.names.multi_lambda.zip
    full_path = '{}/{}'.format(path, zip_name)
    resp = s3.get_object(Bucket=bosslet_config.LAMBDA_BUCKET, Key=zip_name)

    bytes = resp['Body'].read()

    with open(full_path , 'wb') as out:
        out.write(bytes)

    print('Saved zip to {}'.format(full_path))

def upload_lambda_zip(bosslet_config, path):
    """
    Upload a  multilambda.domain.zip to the S3 bucket.  Useful when
    developing and small changes need to be made to a lambda function, but a full
    rebuild of the entire zip file isn't required.
    """
    s3 = bosslet_config.session.client('s3')
    with open(path, 'rb') as in_file:
        resp = s3.put_object(Bucket=bosslet_config.LAMBDA_BUCKET,
                             Key=bosslet_config.names.multi_lambda.zip,
                             Body=in_file)
    print(resp)

<|endoftext|>


<commit_before><commit_msg>1 python-list.py<commit_after><|endoftext|>


<commit_before>"""
==========================
Plotting tract profiles
==========================

An example of tracking and segmenting two tracts, and plotting their tract
profiles for FA (calculated with DTI).

"""
import os.path as op
import matplotlib.pyplot as plt
import numpy as np
import nibabel as nib
import dipy.data as dpd
from dipy.data import fetcher
import dipy.tracking.utils as dtu
import dipy.tracking.streamline as dts
from dipy.io.streamline import save_tractogram

import AFQ.utils.streamlines as aus
import AFQ.data as afd
import AFQ.tractography as aft
import AFQ.registration as reg
import AFQ.dti as dti
import AFQ.segmentation as seg
from AFQ.utils.volume import patch_up_roi

dpd.fetch_stanford_hardi()

hardi_dir = op.join(fetcher.dipy_home, "stanford_hardi")
hardi_fdata = op.join(hardi_dir, "HARDI150.nii.gz")
hardi_fbval = op.join(hardi_dir, "HARDI150.bval")
hardi_fbvec = op.join(hardi_dir, "HARDI150.bvec")

img = nib.load(hardi_fdata)

print("Calculating DTI...")
if not op.exists('./dti_FA.nii.gz'):
    dti_params = dti.fit_dti(hardi_fdata, hardi_fbval, hardi_fbvec,
                             out_dir='.')
else:
    dti_params = {'FA': './dti_FA.nii.gz',
                  'params': './dti_params.nii.gz'}

FA_img = nib.load(dti_params['FA'])
FA_data = FA_img.get_fdata()

templates = afd.read_templates()
bundle_names = ["CST", "ILF"]

bundles = {}
for name in bundle_names:
    for hemi in ['_R', '_L']:
        bundles[name + hemi] = {
            'ROIs': [templates[name + '_roi1' + hemi],
                     templates[name + '_roi2' + hemi]],
            'rules': [True, True],
            'prob_map': templates[name + hemi + '_prob_map'],
            'cross_midline': False}

print("Registering to template...")
MNI_T2_img = dpd.read_mni_template()
if not op.exists('mapping.nii.gz'):
    import dipy.core.gradients as dpg
    gtab = dpg.gradient_table(hardi_fbval, hardi_fbvec)
    warped_hardi, mapping = reg.syn_register_dwi(hardi_fdata, gtab)
    reg.write_mapping(mapping, './mapping.nii.gz')
else:
    mapping = reg.read_mapping('./mapping.nii.gz', img, MNI_T2_img)

print("Tracking...")
if not op.exists('dti_streamlines.trk'):
    seed_roi = np.zeros(img.shape[:-1])
    for name in bundle_names:
        for hemi in ['_R', '_L']:
            for roi in bundles[name + hemi]['ROIs']:
                warped_roi = patch_up_roi(
                    (mapping.transform_inverse(
                        roi.get_data().astype(np.float32),
                     interpolation='linear')) > 0)
                seed_roi += warped_roi

    streamlines = aft.track(dti_params['params'], seed_mask=seed_roi,
                            stop_mask=FA_data > 0.2)

    save_tractogram('./dti_streamlines.trk', streamlines, img.affine)
else:
    tg = nib.streamlines.load('./dti_streamlines.trk').tractogram
    streamlines = tg.streamlines

streamlines = dts.Streamlines(dtu.move_streamlines(
    [s for s in streamlines if s.shape[0] > 100],
    np.linalg.inv(img.affine)))

print("Segmenting fiber groups...")
fiber_groups = seg.segment(hardi_fdata,
                           hardi_fbval,
                           hardi_fbvec,
                           streamlines,
                           bundles,
                           reg_template=MNI_T2_img,
                           mapping=mapping)

print("Cleaning fiber groups...")
for bundle in bundles:
    fiber_groups[bundle] = seg.clean_fiber_group(fiber_groups[bundle])

print("Extracting tract profiles...")
for bundle in bundles:
    fig, ax = plt.subplots(1)
    weights = seg.gaussian_weights(fiber_groups[bundle])
    profile = seg.calculate_tract_profile(FA_data, fiber_groups[bundle],
                                          weights=weights)
    ax.plot(profile)
    ax.set_title(bundle)

plt.show()
<commit_msg>Saving this with the identity matrix now opens well with mi-brain.<commit_after>"""
==========================
Plotting tract profiles
==========================

An example of tracking and segmenting two tracts, and plotting their tract
profiles for FA (calculated with DTI).

"""
import os.path as op
import matplotlib.pyplot as plt
import numpy as np
import nibabel as nib
import dipy.data as dpd
from dipy.data import fetcher
import dipy.tracking.utils as dtu
import dipy.tracking.streamline as dts
from dipy.io.streamline import save_tractogram

import AFQ.utils.streamlines as aus
import AFQ.data as afd
import AFQ.tractography as aft
import AFQ.registration as reg
import AFQ.dti as dti
import AFQ.segmentation as seg
from AFQ.utils.volume import patch_up_roi

dpd.fetch_stanford_hardi()

hardi_dir = op.join(fetcher.dipy_home, "stanford_hardi")
hardi_fdata = op.join(hardi_dir, "HARDI150.nii.gz")
hardi_fbval = op.join(hardi_dir, "HARDI150.bval")
hardi_fbvec = op.join(hardi_dir, "HARDI150.bvec")

img = nib.load(hardi_fdata)

print("Calculating DTI...")
if not op.exists('./dti_FA.nii.gz'):
    dti_params = dti.fit_dti(hardi_fdata, hardi_fbval, hardi_fbvec,
                             out_dir='.')
else:
    dti_params = {'FA': './dti_FA.nii.gz',
                  'params': './dti_params.nii.gz'}

FA_img = nib.load(dti_params['FA'])
FA_data = FA_img.get_fdata()

templates = afd.read_templates()
bundle_names = ["CST", "ILF"]

bundles = {}
for name in bundle_names:
    for hemi in ['_R', '_L']:
        bundles[name + hemi] = {
            'ROIs': [templates[name + '_roi1' + hemi],
                     templates[name + '_roi2' + hemi]],
            'rules': [True, True],
            'prob_map': templates[name + hemi + '_prob_map'],
            'cross_midline': False}

print("Registering to template...")
MNI_T2_img = dpd.read_mni_template()
if not op.exists('mapping.nii.gz'):
    import dipy.core.gradients as dpg
    gtab = dpg.gradient_table(hardi_fbval, hardi_fbvec)
    warped_hardi, mapping = reg.syn_register_dwi(hardi_fdata, gtab)
    reg.write_mapping(mapping, './mapping.nii.gz')
else:
    mapping = reg.read_mapping('./mapping.nii.gz', img, MNI_T2_img)

print("Tracking...")
if not op.exists('dti_streamlines.trk'):
    seed_roi = np.zeros(img.shape[:-1])
    for name in bundle_names:
        for hemi in ['_R', '_L']:
            for roi in bundles[name + hemi]['ROIs']:
                warped_roi = patch_up_roi(
                    (mapping.transform_inverse(
                        roi.get_data().astype(np.float32),
                     interpolation='linear')) > 0)
                seed_roi += warped_roi

    streamlines = aft.track(dti_params['params'], seed_mask=seed_roi,
                            stop_mask=FA_data > 0.2)

    save_tractogram('./dti_streamlines.trk', streamlines, np.eye(4))
else:
    tg = nib.streamlines.load('./dti_streamlines.trk').tractogram
    streamlines = tg.streamlines

streamlines = dts.Streamlines(dtu.move_streamlines(
    [s for s in streamlines if s.shape[0] > 100],
    np.linalg.inv(img.affine)))

print("Segmenting fiber groups...")
fiber_groups = seg.segment(hardi_fdata,
                           hardi_fbval,
                           hardi_fbvec,
                           streamlines,
                           bundles,
                           reg_template=MNI_T2_img,
                           mapping=mapping)

print("Cleaning fiber groups...")
for bundle in bundles:
    fiber_groups[bundle] = seg.clean_fiber_group(fiber_groups[bundle])

print("Extracting tract profiles...")
for bundle in bundles:
    fig, ax = plt.subplots(1)
    weights = seg.gaussian_weights(fiber_groups[bundle])
    profile = seg.calculate_tract_profile(FA_data, fiber_groups[bundle],
                                          weights=weights)
    ax.plot(profile)
    ax.set_title(bundle)

plt.show()
<|endoftext|>


<commit_before>VERSION = '0.6.4'
<commit_msg>Increased version to 0.6.5<commit_after>VERSION = '0.6.5'
<|endoftext|>


<commit_before>from __future__ import division

import numpy as np

from chainercv.datasets import CamVidDataset

n_class = 12
dataset = CamVidDataset(split='train')

n_cls_pixels = [0 for _ in range(n_class)]
n_img_pixels = [0 for _ in range(n_class)]

for img, lbl in dataset:
    for cls_i in np.unique(lbl):
        n_cls_pixels[cls_i] += np.sum(lbl == cls_i)
        n_img_pixels[cls_i] += lbl.size
freq = []
for n_cls_pixel, n_img_pixel in zip(n_cls_pixels, n_img_pixels):
    freq.append(n_cls_pixel / n_img_pixel)
freq = np.array(freq)
median_freq = np.median(freq)

np.save('class_weight', median_freq / freq)
<commit_msg>fix calc_weight<commit_after>from __future__ import division

import numpy as np

from chainercv.datasets import CamVidDataset

n_class = 12
dataset = CamVidDataset(split='train')

n_cls_pixels = np.zeros((n_class,))
n_img_pixels = np.zeros((n_class,))

for img, label in dataset:
    for cls_i in np.unique(label):
        n_cls_pixels[cls_i] += np.sum(label == cls_i)
        n_img_pixels[cls_i] += label.size
freq = n_cls_pixels / n_img_pixels
median_freq = np.median(freq)

np.save('class_weight', median_freq / freq)
<|endoftext|>


<commit_before>from exporters.default_retries import retry_long
from exporters.writers.filebase_base_writer import FilebaseBaseWriter

DEFAULT_BUCKET_REGION = 'us-east-1'

class S3Writer(FilebaseBaseWriter):
    """
    Writes items to S3 bucket.

        - bucket (str)
            Name of the bucket to write the items to.

        - aws_access_key_id (str)
            Public acces key to the s3 bucket.

        - aws_secret_access_key (str)
            Secret access key to the s3 bucket.

        - filebase (str)
            Base path to store the items in the bucket.

        - aws_region (str)
            AWS region to connect to.

        - save_metadata (bool)
            Save key's items count as metadata. Default: True
    """
    supported_options = {
        'bucket': {'type': basestring},
        'aws_access_key_id': {'type': basestring, 'env_fallback': 'EXPORTERS_S3WRITER_AWS_LOGIN'},
        'aws_secret_access_key': {'type': basestring, 'env_fallback': 'EXPORTERS_S3WRITER_AWS_SECRET'},
        'aws_region': {'type': basestring, 'default': None},
        'save_metadata': {'type': bool, 'default': True, 'required': False}
    }

    def __init__(self, options):
        import boto

        super(S3Writer, self).__init__(options)
        access_key = self.read_option('aws_access_key_id')
        secret_key = self.read_option('aws_secret_access_key')
        self.aws_region = self.read_option('aws_region')
        bucket_name = self.read_option('bucket')

        if self.aws_region is None:
            try:
                self.aws_region = self._get_bucket_location(access_key, secret_key,
                                                            bucket_name)
            except:
                self.aws_region = DEFAULT_BUCKET_REGION

        self.conn = boto.s3.connect_to_region(self.aws_region,
                                              aws_access_key_id=access_key,
                                              aws_secret_access_key=secret_key)
        self.bucket = self.conn.get_bucket(bucket_name)
        self.save_metadata = self.read_option('save_metadata')
        self.logger.info('S3Writer has been initiated.'
                         'Writing to s3://{}/{}'.format(self.bucket.name, self.filebase))

    def _get_bucket_location(self, access_key, secret_key, bucket):
        import boto

        return boto.connect_s3(access_key, secret_key).get_bucket(bucket).get_location() or DEFAULT_BUCKET_REGION

    @retry_long
    def _write_s3_key(self, dump_path, key_name):
        destination = 's3://{}/{}'.format(self.bucket.name, key_name)
        self.logger.info('Start uploading {} to {}'.format(dump_path, destination))
        key = self.bucket.new_key(key_name)
        key.set_metadata('total', self.write_buffer.stats['written_keys']['keys'].values()[0]['number_of_records'])
        with open(dump_path, 'r') as f:
            key.set_contents_from_file(f)
        key.close()

        self.logger.info('Saved {}'.format(destination))

    def write(self, dump_path, group_key=None):
        if group_key is None:
            group_key = []

        filebase_path, filename = self.create_filebase_name(group_key)
        key_name = filebase_path + '/' + filename
        self._write_s3_key(dump_path, key_name)
<commit_msg>metadata saving fixed<commit_after>from contextlib import closing
from exporters.default_retries import retry_long
from exporters.writers.filebase_base_writer import FilebaseBaseWriter

DEFAULT_BUCKET_REGION = 'us-east-1'

class S3Writer(FilebaseBaseWriter):
    """
    Writes items to S3 bucket.

        - bucket (str)
            Name of the bucket to write the items to.

        - aws_access_key_id (str)
            Public acces key to the s3 bucket.

        - aws_secret_access_key (str)
            Secret access key to the s3 bucket.

        - filebase (str)
            Base path to store the items in the bucket.

        - aws_region (str)
            AWS region to connect to.

        - save_metadata (bool)
            Save key's items count as metadata. Default: True
    """
    supported_options = {
        'bucket': {'type': basestring},
        'aws_access_key_id': {'type': basestring, 'env_fallback': 'EXPORTERS_S3WRITER_AWS_LOGIN'},
        'aws_secret_access_key': {'type': basestring, 'env_fallback': 'EXPORTERS_S3WRITER_AWS_SECRET'},
        'aws_region': {'type': basestring, 'default': None},
        'save_metadata': {'type': bool, 'default': True, 'required': False}
    }

    def __init__(self, options):
        import boto

        super(S3Writer, self).__init__(options)
        access_key = self.read_option('aws_access_key_id')
        secret_key = self.read_option('aws_secret_access_key')
        self.aws_region = self.read_option('aws_region')
        bucket_name = self.read_option('bucket')

        if self.aws_region is None:
            try:
                self.aws_region = self._get_bucket_location(access_key, secret_key,
                                                            bucket_name)
            except:
                self.aws_region = DEFAULT_BUCKET_REGION

        self.conn = boto.s3.connect_to_region(self.aws_region,
                                              aws_access_key_id=access_key,
                                              aws_secret_access_key=secret_key)
        self.bucket = self.conn.get_bucket(bucket_name)
        self.save_metadata = self.read_option('save_metadata')
        self.logger.info('S3Writer has been initiated.'
                         'Writing to s3://{}/{}'.format(self.bucket.name, self.filebase))

    def _get_bucket_location(self, access_key, secret_key, bucket):
        import boto

        return boto.connect_s3(access_key, secret_key).get_bucket(bucket).get_location() or DEFAULT_BUCKET_REGION

    def _get_total_count(self, dump_path):
        return self.write_buffer.stats['written_keys']['keys'][dump_path]['number_of_records']

    @retry_long
    def _write_s3_key(self, dump_path, key_name):
        destination = 's3://{}/{}'.format(self.bucket.name, key_name)
        self.logger.info('Start uploading {} to {}'.format(dump_path, destination))
        with closing(self.bucket.new_key(key_name)) as key, open(dump_path, 'r') as f:
            if self.save_metadata:
                key.set_metadata('total', self._get_total_count(dump_path))
            key.set_contents_from_file(f)

        self.logger.info('Saved {}'.format(destination))

    def write(self, dump_path, group_key=None):
        if group_key is None:
            group_key = []

        filebase_path, filename = self.create_filebase_name(group_key)
        key_name = filebase_path + '/' + filename
        self._write_s3_key(dump_path, key_name)
<|endoftext|>


<commit_before>"""
Base AbuseCH feed handler. The specific AbuseCH RSS feeds bots
(ZeuS Tracker etc.) are built on this.

Maintainer: Jussi Eronen <exec@iki.fi>
"""

import re
import cgi
import time
import socket
import urlparse

from abusehelper.core import events
from abusehelper.contrib.rssbot.rssbot import RSSBot

def is_ip(string):
    for addr_type in (socket.AF_INET, socket.AF_INET6):
        try:
            socket.inet_pton(addr_type, string)
        except (ValueError, socket.error):
            pass
        else:
            return True
    return False

def parse_title(title):
    """
    ZeuS Tracker and SpyEye Tracker styles:

    >>> list(parse_title("1.2.3.4/badness.php (2012-01-01)"))
    [('time', '2012-01-01 00:00:00 UTC')]

    >>> list(parse_title("1.2.3.4/badness.php (2012-01-01 01:02:03)"))
    [('time', '2012-01-01 01:02:03 UTC')]

    Palevo tracker style:

    >>> list(parse_title("1.2.3.4/badness.php 2012-01-01"))
    [('time', '2012-01-01 00:00:00 UTC')]

    Ignore if can't be parsed:

    >>> list(parse_title("1.2.3.4/badness.php"))
    []
    """

    match = re.search(r"(\d{4}-\d\d-\d\d(?: \d\d\:\d\d:\d\d)?)\)?\s*$", title)
    if not match:
        return

    timestamp = match.group(1)
    for format in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
        try:
            timestamp = time.strptime(timestamp, format)
        except ValueError:
            continue
        yield "time", time.strftime("%Y-%m-%d %H:%M:%S UTC", timestamp)
        break

def parse_link(link):
    """
    >>> sorted(parse_link("https://spyeyetracker.abuse.ch/monitor.php?host=www.example.com"))
    [('host', 'www.example.com'), ('more info', 'https://spyeyetracker.abuse.ch/monitor.php?host=www.example.com')]

    >>> sorted(parse_link("https://spyeyetracker.abuse.ch/monitor.php?host=1.2.3.4"))
    [('host', '1.2.3.4'), ('ip', '1.2.3.4'), ('more info', 'https://spyeyetracker.abuse.ch/monitor.php?host=1.2.3.4')]
    """

    yield "more info", link

    parsed = urlparse.urlparse(link)
    query = cgi.parse_qs(parsed[4])
    for host in query.get("host", []):
        yield "host", host

        if is_ip(host):
            yield "ip", host

_levels = {
    "1": "bulletproof hosted",
    "2": "hacked webserver",
    "3": "free hosting service",
    "4": "unknown",
    "5": "hosted on a fastflux botnet"
}

_sbl_prefix = "http://www.spamhaus.org/sbl/sbl.lasso?query="

def parse_description(description):
    for part in description.split(","):
        pair = part.split(":", 1)
        if len(pair) < 2:
            continue
        key, value = pair

        value = value.strip()
        if not value:
            continue

        key = key.strip().lower()
        if key == "as":
            if value.lower().startswith("as"):
                value = value[2:]
            yield "asn", value
        elif key == "ip address":
            yield "ip", value
        elif key == "sbl":
            if value.lower() != "not listed":
                yield "sbl", _sbl_prefix + value
        elif key == "level":
            yield "level", _levels.get(value, value)
        else:
            yield key, value

class AbuseCHBot(RSSBot):
    feeds = None

    malwares = {
        "zeus": {
            "c&c": "https://zeustracker.abuse.ch/rss.php",
            "config": "https://zeustracker.abuse.ch/monitor.php?urlfeed=configs",
            "binary": "https://zeustracker.abuse.ch/monitor.php?urlfeed=binaries",
            "dropzone": "https://zeustracker.abuse.ch/monitor.php?urlfeed=dropzones"
        },
        "spyeye": {
            "c&c": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=tracker",
            "config": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=configurls",
            "binary": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=binaryurls",
            "dropzone": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=dropurls"
        },
        "palevo": {
            "c&c": "http://amada.abuse.ch/palevotracker.php?rssfeed"
        }
    }

    @property
    def urls(self):
        if self._urls is not None:
            return self._urls

        self._urls = dict()
        for malware, feed_types in self.malwares.iteritems():
            for feed_type, url in feed_types.iteritems():
                self._urls[url] = malware, feed_type
        return self._urls
    _urls = None

    def feed_keys(self, malwares=None, types=None, **_):
        if malwares is None:
            malwares = self.malwares

        for malware in malwares:
            if not malware in self.malwares:
                self.log.error("no support for malware %r" % malware)
                continue

            feed_types = self.malwares[malware]
            if types is None:
                types = feed_types

            for feed_type in types:
                if feed_type not in feed_types:
                    self.log.error("no support for %r type %r" % (malware, feed_type))
                    continue

                yield (feed_types[feed_type],)

    def parse_title(self, string):
        return parse_title(string)

    def parse_link(self, string):
        return parse_link(string)

    def parse_description(self, string):
        return parse_description(string)

    def create_event(self, source=None, **keys):
        event = events.Event()

        for name, string in keys.iteritems():
            parse = getattr(self, "parse_" + name, None)
            if parse is None:
                continue

            for key, value in parse(string):
                event.add(key, value)

        if not event.contains():
            return None

        if source in self.urls:
            malware, feed_type = self.urls[source]
            event.add("malware", malware)
            event.add("type", feed_type)
        event.add("source url", source)

        return event

if __name__ == "__main__":
    AbuseCHBot.from_command_line().run()
<commit_msg>abusehelper.contrib.abusech.abusechbot: Updated new palevotracker url.<commit_after>"""
Base AbuseCH feed handler. The specific AbuseCH RSS feeds bots
(ZeuS Tracker etc.) are built on this.

Maintainer: Jussi Eronen <exec@iki.fi>
"""

import re
import cgi
import time
import socket
import urlparse

from abusehelper.core import events
from abusehelper.contrib.rssbot.rssbot import RSSBot

def is_ip(string):
    for addr_type in (socket.AF_INET, socket.AF_INET6):
        try:
            socket.inet_pton(addr_type, string)
        except (ValueError, socket.error):
            pass
        else:
            return True
    return False

def parse_title(title):
    """
    ZeuS Tracker and SpyEye Tracker styles:

    >>> list(parse_title("1.2.3.4/badness.php (2012-01-01)"))
    [('time', '2012-01-01 00:00:00 UTC')]

    >>> list(parse_title("1.2.3.4/badness.php (2012-01-01 01:02:03)"))
    [('time', '2012-01-01 01:02:03 UTC')]

    Palevo tracker style:

    >>> list(parse_title("1.2.3.4/badness.php 2012-01-01"))
    [('time', '2012-01-01 00:00:00 UTC')]

    Ignore if can't be parsed:

    >>> list(parse_title("1.2.3.4/badness.php"))
    []
    """

    match = re.search(r"(\d{4}-\d\d-\d\d(?: \d\d\:\d\d:\d\d)?)\)?\s*$", title)
    if not match:
        return

    timestamp = match.group(1)
    for format in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
        try:
            timestamp = time.strptime(timestamp, format)
        except ValueError:
            continue
        yield "time", time.strftime("%Y-%m-%d %H:%M:%S UTC", timestamp)
        break

def parse_link(link):
    """
    >>> sorted(parse_link("https://spyeyetracker.abuse.ch/monitor.php?host=www.example.com"))
    [('host', 'www.example.com'), ('more info', 'https://spyeyetracker.abuse.ch/monitor.php?host=www.example.com')]

    >>> sorted(parse_link("https://spyeyetracker.abuse.ch/monitor.php?host=1.2.3.4"))
    [('host', '1.2.3.4'), ('ip', '1.2.3.4'), ('more info', 'https://spyeyetracker.abuse.ch/monitor.php?host=1.2.3.4')]
    """

    yield "more info", link

    parsed = urlparse.urlparse(link)
    query = cgi.parse_qs(parsed[4])
    for host in query.get("host", []):
        yield "host", host

        if is_ip(host):
            yield "ip", host

_levels = {
    "1": "bulletproof hosted",
    "2": "hacked webserver",
    "3": "free hosting service",
    "4": "unknown",
    "5": "hosted on a fastflux botnet"
}

_sbl_prefix = "http://www.spamhaus.org/sbl/sbl.lasso?query="

def parse_description(description):
    for part in description.split(","):
        pair = part.split(":", 1)
        if len(pair) < 2:
            continue
        key, value = pair

        value = value.strip()
        if not value:
            continue

        key = key.strip().lower()
        if key == "as":
            if value.lower().startswith("as"):
                value = value[2:]
            yield "asn", value
        elif key == "ip address":
            yield "ip", value
        elif key == "sbl":
            if value.lower() != "not listed":
                yield "sbl", _sbl_prefix + value
        elif key == "level":
            yield "level", _levels.get(value, value)
        else:
            yield key, value

class AbuseCHBot(RSSBot):
    feeds = None

    malwares = {
        "zeus": {
            "c&c": "https://zeustracker.abuse.ch/rss.php",
            "config": "https://zeustracker.abuse.ch/monitor.php?urlfeed=configs",
            "binary": "https://zeustracker.abuse.ch/monitor.php?urlfeed=binaries",
            "dropzone": "https://zeustracker.abuse.ch/monitor.php?urlfeed=dropzones"
        },
        "spyeye": {
            "c&c": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=tracker",
            "config": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=configurls",
            "binary": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=binaryurls",
            "dropzone": "https://spyeyetracker.abuse.ch/monitor.php?rssfeed=dropurls"
        },
        "palevo": {
            "c&c": "https://palevotracker.abuse.ch/?rssfeed"
        }
    }

    @property
    def urls(self):
        if self._urls is not None:
            return self._urls

        self._urls = dict()
        for malware, feed_types in self.malwares.iteritems():
            for feed_type, url in feed_types.iteritems():
                self._urls[url] = malware, feed_type
        return self._urls
    _urls = None

    def feed_keys(self, malwares=None, types=None, **_):
        if malwares is None:
            malwares = self.malwares

        for malware in malwares:
            if not malware in self.malwares:
                self.log.error("no support for malware %r" % malware)
                continue

            feed_types = self.malwares[malware]
            if types is None:
                types = feed_types

            for feed_type in types:
                if feed_type not in feed_types:
                    self.log.error("no support for %r type %r" % (malware, feed_type))
                    continue

                yield (feed_types[feed_type],)

    def parse_title(self, string):
        return parse_title(string)

    def parse_link(self, string):
        return parse_link(string)

    def parse_description(self, string):
        return parse_description(string)

    def create_event(self, source=None, **keys):
        event = events.Event()

        for name, string in keys.iteritems():
            parse = getattr(self, "parse_" + name, None)
            if parse is None:
                continue

            for key, value in parse(string):
                event.add(key, value)

        if not event.contains():
            return None

        if source in self.urls:
            malware, feed_type = self.urls[source]
            event.add("malware", malware)
            event.add("type", feed_type)
        event.add("source url", source)

        return event

if __name__ == "__main__":
    AbuseCHBot.from_command_line().run()
<|endoftext|>


<commit_before>import socket
import idiokit
import urlparse
from abusehelper.core import utils, bot, events

FEED_URL = "http://vxvault.siri-urz.net/URL_List.php"

def i_am_a_name(string):
    for addr_type in (socket.AF_INET, socket.AF_INET6):
        try:
            socket.inet_ntop(addr_type, socket.inet_pton(addr_type, string))
        except (ValueError, socket.error):
            pass
        else:
            return False
    return True

def parseURL(line):
    parsed = urlparse.urlparse(line)
    if parsed.scheme and parsed.netloc:
        return line, parsed.netloc
    else:
        return None, None

class VxVaultBot(bot.PollingBot):
    feed_url = bot.Param(default=FEED_URL)

    def poll(self):
        pipe = self._poll(url=self.feed_url)
        return pipe

    @idiokit.stream
    def _poll(self, url):
        self.log.info("Downloading %s" % url)
        try:
            info, fileobj = yield utils.fetch_url(url)
        except utils.FetchUrlFailed, fuf:
            raise bot.PollSkipped("failed to download {0!r} ({1})".format(url, fuf))
        self.log.info("Downloaded")

        for line in fileobj:
            url, netloc = parseURL(line)
            if url is None:
                continue
            event = events.Event()
            event.add("url", url)
            if i_am_a_name(netloc):
                event.add("domain name", netloc)
            else:
                event.add("ip", netloc)
            event.add("feed", "vxvault")
            event.add("feed url", self.feed_url)
            event.add("type", "malware")
            event.add("description", "This host is most likely hosting a malware URL.")
            yield idiokit.send(event)

if __name__ == "__main__":
    VxVaultBot.from_command_line().execute()

<commit_msg>abusehelper.contrib.vxvault.vxvaultbot: Remove indirection.<commit_after>import socket
import idiokit
import urlparse
from abusehelper.core import utils, bot, events

FEED_URL = "http://vxvault.siri-urz.net/URL_List.php"

def i_am_a_name(string):
    for addr_type in (socket.AF_INET, socket.AF_INET6):
        try:
            socket.inet_ntop(addr_type, socket.inet_pton(addr_type, string))
        except (ValueError, socket.error):
            pass
        else:
            return False
    return True

def parseURL(line):
    parsed = urlparse.urlparse(line)
    if parsed.scheme and parsed.netloc:
        return line, parsed.netloc
    else:
        return None, None

class VxVaultBot(bot.PollingBot):
    feed_url = bot.Param(default=FEED_URL)

    @idiokit.stream
    def poll(self):
        self.log.info("Downloading %s" % self.feed_url)
        try:
            info, fileobj = yield utils.fetch_url(self.feed_url)
        except utils.FetchUrlFailed, fuf:
            raise bot.PollSkipped("failed to download {0!r} ({1})".format(self.feed_url, fuf))
        self.log.info("Downloaded")

        for line in fileobj:
            url, netloc = parseURL(line)
            if url is None:
                continue
            event = events.Event()
            event.add("url", url)
            if i_am_a_name(netloc):
                event.add("domain name", netloc)
            else:
                event.add("ip", netloc)
            event.add("feed", "vxvault")
            event.add("feed url", self.feed_url)
            event.add("type", "malware")
            event.add("description", "This host is most likely hosting a malware URL.")
            yield idiokit.send(event)

if __name__ == "__main__":
    VxVaultBot.from_command_line().execute()

<|endoftext|>


<commit_before>from __future__ import division
# LIBTBX_SET_DISPATCHER_NAME dev.xia2.rogues_gallery

import exceptions
import json
import os
import sys

# Needed to make xia2 imports work correctly FIXME we should make these go away
import libtbx.load_env
xia2_root_dir = libtbx.env.find_in_repositories("xia2", optional=False)
sys.path.insert(0, xia2_root_dir)
os.environ['XIA2_ROOT'] = xia2_root_dir
os.environ['XIA2CORE_ROOT'] = os.path.join(xia2_root_dir, "core")

def munch_rogues(rogues):
  rogue_reflections = []
  for record in open(rogues):
    if not record.strip():
      continue
    tokens = record.split()
    if not tokens[-1] == '*':
      continue
    x = float(tokens[15])
    y = float(tokens[16])
    b = int(tokens[6])
    h, k, l = map(int, tokens[3:6]) # don't forget these are probably reindexed
    rogue_reflections.append((b, x, y, h, k, l))

  return rogue_reflections

def reconstruct_rogues():
  assert os.path.exists('xia2.json')
  from Schema.XProject import XProject
  xinfo = XProject.from_json(filename='xia2.json')

  from dxtbx.model.experiment.experiment_list import ExperimentListFactory
  import cPickle as pickle
  import dials # because WARNING:root:No profile class gaussian_rs registered
  crystals = xinfo.get_crystals()
  assert len(crystals) == 1

  for xname in crystals:
    crystal = crystals[xname]

  scaler = crystal._get_scaler()

  epochs = scaler._sweep_handler.get_epochs()

  rogues = os.path.join(scaler.get_working_directory(),
                        xname, 'scale', 'ROGUES')

  rogue_reflections = munch_rogues(rogues)

  batched_reflections = { }

  for epoch in epochs:
    si = scaler._sweep_handler.get_sweep_information(epoch)
    intgr = si.get_integrater()
    experiments = ExperimentListFactory.from_json_file(
      intgr.get_integrated_experiments())
    reflections = pickle.load(open(intgr.get_integrated_reflections()))
    batched_reflections[si.get_batch_range()] = (experiments, reflections)

if __name__ == '__main__':
  reconstruct_rogues()
<commit_msg>Now have code which will find nearest matching reflection using annlib; seems to work; get the bb; now need to work out how to get the pixel data out<commit_after>from __future__ import division
# LIBTBX_SET_DISPATCHER_NAME dev.xia2.rogues_gallery

import exceptions
import json
import os
import sys

# Needed to make xia2 imports work correctly FIXME we should make these go away
import libtbx.load_env
xia2_root_dir = libtbx.env.find_in_repositories("xia2", optional=False)
sys.path.insert(0, xia2_root_dir)
os.environ['XIA2_ROOT'] = xia2_root_dir
os.environ['XIA2CORE_ROOT'] = os.path.join(xia2_root_dir, "core")

def munch_rogues(rogues):
  rogue_reflections = []
  for record in open(rogues):
    if not record.strip():
      continue
    tokens = record.split()
    if not tokens[-1] == '*':
      continue
    x, y, z = map(float, tokens[15:18])
    b = int(tokens[6])
    h, k, l = map(int, tokens[3:6]) # don't forget these are probably reindexed
    rogue_reflections.append((b, x, y, z, h, k, l))

  return rogue_reflections

def reconstruct_rogues():
  assert os.path.exists('xia2.json')
  from Schema.XProject import XProject
  xinfo = XProject.from_json(filename='xia2.json')

  from dxtbx.model.experiment.experiment_list import ExperimentListFactory
  import cPickle as pickle
  import dials # because WARNING:root:No profile class gaussian_rs registered
  crystals = xinfo.get_crystals()
  assert len(crystals) == 1

  for xname in crystals:
    crystal = crystals[xname]

  scaler = crystal._get_scaler()

  epochs = scaler._sweep_handler.get_epochs()

  rogues = os.path.join(scaler.get_working_directory(),
                        xname, 'scale', 'ROGUES')

  rogue_reflections = munch_rogues(rogues)

  batched_reflections = { }

  for epoch in epochs:
    si = scaler._sweep_handler.get_sweep_information(epoch)
    intgr = si.get_integrater()
    experiments = ExperimentListFactory.from_json_file(
      intgr.get_integrated_experiments())
    reflections = pickle.load(open(intgr.get_integrated_reflections()))
    batched_reflections[si.get_batch_range()] = (experiments, reflections)

  # - look up reflection in reflection list, get bounding box
  # - pull pixels given from image set, flatten these, write out

  from dials.array_family import flex
  from annlib_ext import AnnAdaptor as ann_adaptor

  reflections_run = { }
  for run in batched_reflections:
    reflections_run[run] = []

  for rogue in rogue_reflections:
    b = rogue[0]
    for run in batched_reflections:
      if b >= run[0] and b <= run[1]:
        reflections_run[run].append(rogue)
        break

  for run in reflections_run:
    experiment = batched_reflections[run][0]
    reflections = batched_reflections[run][1]
    rogues = reflections_run[run]
    reference = flex.double()
    scan = experiment.scans()[0]
    images = experiment.imagesets()[0]
    for xyz in reflections['xyzcal.px']:
      reference.append(xyz[0])
      reference.append(xyz[1])
      reference.append(xyz[2])

    search = flex.double()
    for rogue in rogues:
      search.append(rogue[1])
      search.append(rogue[2])
      search.append(scan.get_array_index_from_angle(rogue[3]))

    ann = ann_adaptor(data=reference, dim=3, k=1)
    ann.query(search)

    for j, rogue in enumerate(rogues):
      nn = reflections[ann.nn[j]]
      print nn['bbox']
      # FIXME now extract the pixel information

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


<commit_before># Plugin (install, enable, disable).
# Upgrade all.
# Uninstall.
import logging
from os import path

log = logging.getLogger(__name__)

def __virtual__():
  return "netbeans.exceptions" in __salt__

def _check_version(version):
  """Checks a version for validiti and converts latest to an actual version.

  version:
    The version to check.
  """
  # Update cache if needed.
  if _check_version._versions_cache is None:
    log.debug("Loading versions cache ...")
    _check_version._versions_cache = __salt__["netbeans.list_versions"]()

  # Convert latest.
  if version is None or version == "latest":
    return __salt__["netbeans.pick_latest_version"](
        _check_version._versions_cache
    )

  # Check if version is available.
  if version not in _check_version._versions_cache:
    return None
  return version

# Module level version cache.
_check_version._versions_cache = None

def install(name, version, features='', source_url=None):
  """Installs the named version of NetBeans on the system.

  name:
    The name of the task, not currently used.

  version:
    The version of NetbBeans to install.
    Currently supports only versions 7.3 and above.
    This is due to the naming of the install file in versions below 7.3.

  features:
    Determines which installer to download based on features.
    If left empty (which is the default) the installer for all features
    will be downloaded.
    Otherwise the options are:
      - cpp
      - javaee
      - javase
      - php

  source_url:
    Remote URL from which the installer will be downloaded.
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  resolved_version = _check_version(version)
  if resolved_version is not None:
    version = resolved_version

  # Look for existing installation.
  # Look even for invalid versions because it could be a
  # locally installed version that is not available on the
  # specified source.
  exceptions = __salt__["netbeans.exceptions"]()
  try:
    found = __salt__["netbeans.find_installation"](version)
    result["comment"] = "NetBeans {ver} already installed at {path}.".format(
        path=found, ver=version
    )
    result["result"]  = True
    return result

  except exceptions["NoInstallFound"]:
    pass

  # Check if the version is valid.
  if resolved_version is None:
    result["comment"] = "Invalid version {ver}.".format(ver=version)
    result["result"]  = False
    return result

  # Download installer.
  installer = None
  try:
    installer = __salt__["netbeans.download_installer"](
        version, features=features, url=source_url
    )

  except Exception as ex:
    result["comment"] = "Could not fetch installer: {0}".format(ex)
    result["result"]  = False
    return result

  # Bail now on test.
  if __opts__["test"]:
    result["comment"] = "Installer downloaded to {path}.".format(path=installer)
    result["result"]  = None
    result["changes"] = {"new": version}
    return result

  # Install in silent mode.
  install_code = __salt__["cmd.retcode"](installer + " --silent")
  if install_code != 0:
    result["comment"] = "Installation failed with code {}".format(install_code)
    result["result"]  = False
    return result

  # Find the installation path again.
  install_path = __salt__["netbeans.find_installation"](version)
  result["comment"] = "Installed version {ver} to {path}.".format(
      path=install_path, ver=version
  )
  result["result"]  = True
  result["changes"] = {"new": version}
  return result

def pinstall(name=None, version=None):
  """Install and enable a plugin.

  name:
    The name of the plugin to install.

  version:
    The version of NetBeans to install the plugin for.
    Defaults to latest available.
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  version = _check_version(version)
  if version is None:
    result["comment"] = "Invalid version {} specified.".format(version)
    result["result"]  = False
    return result

  # Look for the plugin state.
  exceptions   = __salt__["netbeans.exceptions"]()
  plugin_state = None
  plugin_version = None
  try:
    (state, version) = __salt__["netbeans.find_plugin"](version, name)
    plugin_state   = state
    plugin_version = version

  except exceptions["NoInstallFound"]:
    result["comment"] = "NetBeans {} not installed.".format(version)
    result["result"]  = False
    return result

  except exceptions["NoPluginFound"]:
    result["comment"] = "Plugin {} not found for NetBeans {}.".format(
        name, version
    )
    result["result"] = False
    return result

  # If up to date and installed done.
  if plugin_state == "enabled":
    result["comment"] = "Plugin already installed and enabled."
    result["result"]  = True

  # Get ready to run the command.
  args    = []
  comment = None
  fail_comment = None
  test_comment = None

  # If not enabled prepare for enable.
  if plugin_state == "installed":
    args.extend(["--enable", name])
    comment = "Plugin enabled."
    fail_comment = "Could not enable plugin."
    test_comment = "Plugin will be enabled."
    result["changes"] = { "state": {
      "new": "enabled",
      "old": "installed"
    } }

  # If not up to date prepare for update.
  if plugin_state == "update":
    args.extend(["--update", name])
    comment = "Plugin updated."
    fail_comment = "Could not update plugin."
    test_comment = "Plugin will be updated."
    result["changes"] = { "version": {
      "new": plugin_version[1],
      "old": plugin_version[0]
    } }

  # If not installed prepare for install.
  if plugin_state == "available":
    args.extend(["--install", name])
    comment = "Plugin installed."
    fail_comment = "Could not install plugin."
    test_comment = "Plugin will be installed."
    result["changes"] = {
      "state": {
        "new": "enabled",
        "old": "available"
      },
      "version": {
        "new": plugin_version,
        "old": None
      }
    }

  if not args:
    result["comment"] = "Unrecognised state {} for plugin.".format(
        plugin_state
    )
    result["result"]  = False
    return result

  # If __opts__["test"] return now.
  if __opts__["test"]:
    result["comment"] = test_comment
    result["result"]  = None
    return result

  # Run the command.
  full_args = [
    "--locale", "en",
    "--nogui",
    "--nosplash",
    "--modules"
  ]
  full_args.extend(args)
  output = __salt__["cmd.retcode"](bin + " " + " ".join(full_args))

def start(name, version=None, args=None):
  """Ensures that NetBeans is running.
  The command line interface is very limited and has the
  huge problem that it actually starts the IDE even when
  module commands are specified.

  If the IDE **is not** running, when the command is completed,
  the process will stay alive in full IDE mode.
  If the IDE **is** already running, when the command is completed,
  the process will terminate.

  This state should be used as a prerequisite of all other
  states so that the IDE is guaranteed to be running.

  version:
    The version of the IDE to run.

  args:
    List of arguments to pass to the command line.
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  version = _check_version(version)
  if version is None:
    result["comment"] = "Invalid version {} specified.".format(version)
    result["result"]  = False
    return result

  # Look for the plugin state.
  exceptions   = __salt__["netbeans.exceptions"]()
  try:
    nb_path = __salt__["netbeans.find_installation"](version)
    bin  = path.join(nb_path, "bin", "netbeans")
    args = ["--locale", "en", "--nogui", "--nosplash"]
    code = __salt__["cmd.shell"](bin + " " + " ".join(args) + " &")

    if code == 0:
      result["comment"] = "Sarted NetBeans in silent mode."
      result["result"]  = True

    else:
      result["comment"] = "Could not sart NetBeans in silent mode."
      result["result"]  = False
    return result

  except exceptions["NoInstallFound"]:
    result["comment"] = "NetBeans {} not installed.".format(version)
    result["result"]  = False
    return result

def stop(name, version=None):
  """
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  version = _check_version(version)
  if version is None:
    result["comment"] = "Invalid version {} specified.".format(version)
    result["result"]  = False
    return result

  # Look for the plugin state.
  exceptions   = __salt__["netbeans.exceptions"]()
  try:
    nb_path = __salt__["netbeans.find_installation"](version)
    native  = path.join(nb_path, "platform", "lib")
    cmd = (
        "ps -ef | grep '" + native +
        "' | grep -v grep | awk '{ print $2 }' | xargs kill"
    )
    __salt__["cmd.shell"](cmd)

  except exceptions["NoInstallFound"]:
    result["comment"] = "NetBeans {} not installed.".format(version)
    result["result"]  = False
    return result
<commit_msg>Using nohup to start netbeans<commit_after># Plugin (install, enable, disable).
# Upgrade all.
# Uninstall.
import logging
from os import path

log = logging.getLogger(__name__)

def __virtual__():
  return "netbeans.exceptions" in __salt__

def _check_version(version):
  """Checks a version for validiti and converts latest to an actual version.

  version:
    The version to check.
  """
  # Update cache if needed.
  if _check_version._versions_cache is None:
    log.debug("Loading versions cache ...")
    _check_version._versions_cache = __salt__["netbeans.list_versions"]()

  # Convert latest.
  if version is None or version == "latest":
    return __salt__["netbeans.pick_latest_version"](
        _check_version._versions_cache
    )

  # Check if version is available.
  if version not in _check_version._versions_cache:
    return None
  return version

# Module level version cache.
_check_version._versions_cache = None

def install(name, version, features='', source_url=None):
  """Installs the named version of NetBeans on the system.

  name:
    The name of the task, not currently used.

  version:
    The version of NetbBeans to install.
    Currently supports only versions 7.3 and above.
    This is due to the naming of the install file in versions below 7.3.

  features:
    Determines which installer to download based on features.
    If left empty (which is the default) the installer for all features
    will be downloaded.
    Otherwise the options are:
      - cpp
      - javaee
      - javase
      - php

  source_url:
    Remote URL from which the installer will be downloaded.
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  resolved_version = _check_version(version)
  if resolved_version is not None:
    version = resolved_version

  # Look for existing installation.
  # Look even for invalid versions because it could be a
  # locally installed version that is not available on the
  # specified source.
  exceptions = __salt__["netbeans.exceptions"]()
  try:
    found = __salt__["netbeans.find_installation"](version)
    result["comment"] = "NetBeans {ver} already installed at {path}.".format(
        path=found, ver=version
    )
    result["result"]  = True
    return result

  except exceptions["NoInstallFound"]:
    pass

  # Check if the version is valid.
  if resolved_version is None:
    result["comment"] = "Invalid version {ver}.".format(ver=version)
    result["result"]  = False
    return result

  # Download installer.
  installer = None
  try:
    installer = __salt__["netbeans.download_installer"](
        version, features=features, url=source_url
    )

  except Exception as ex:
    result["comment"] = "Could not fetch installer: {0}".format(ex)
    result["result"]  = False
    return result

  # Bail now on test.
  if __opts__["test"]:
    result["comment"] = "Installer downloaded to {path}.".format(path=installer)
    result["result"]  = None
    result["changes"] = {"new": version}
    return result

  # Install in silent mode.
  install_code = __salt__["cmd.retcode"](installer + " --silent")
  if install_code != 0:
    result["comment"] = "Installation failed with code {}".format(install_code)
    result["result"]  = False
    return result

  # Find the installation path again.
  install_path = __salt__["netbeans.find_installation"](version)
  result["comment"] = "Installed version {ver} to {path}.".format(
      path=install_path, ver=version
  )
  result["result"]  = True
  result["changes"] = {"new": version}
  return result

def pinstall(name=None, version=None):
  """Install and enable a plugin.

  name:
    The name of the plugin to install.

  version:
    The version of NetBeans to install the plugin for.
    Defaults to latest available.
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  version = _check_version(version)
  if version is None:
    result["comment"] = "Invalid version {} specified.".format(version)
    result["result"]  = False
    return result

  # Look for the plugin state.
  exceptions   = __salt__["netbeans.exceptions"]()
  plugin_state = None
  plugin_version = None
  try:
    (state, version) = __salt__["netbeans.find_plugin"](version, name)
    plugin_state   = state
    plugin_version = version

  except exceptions["NoInstallFound"]:
    result["comment"] = "NetBeans {} not installed.".format(version)
    result["result"]  = False
    return result

  except exceptions["NoPluginFound"]:
    result["comment"] = "Plugin {} not found for NetBeans {}.".format(
        name, version
    )
    result["result"] = False
    return result

  # If up to date and installed done.
  if plugin_state == "enabled":
    result["comment"] = "Plugin already installed and enabled."
    result["result"]  = True

  # Get ready to run the command.
  args    = []
  comment = None
  fail_comment = None
  test_comment = None

  # If not enabled prepare for enable.
  if plugin_state == "installed":
    args.extend(["--enable", name])
    comment = "Plugin enabled."
    fail_comment = "Could not enable plugin."
    test_comment = "Plugin will be enabled."
    result["changes"] = { "state": {
      "new": "enabled",
      "old": "installed"
    } }

  # If not up to date prepare for update.
  if plugin_state == "update":
    args.extend(["--update", name])
    comment = "Plugin updated."
    fail_comment = "Could not update plugin."
    test_comment = "Plugin will be updated."
    result["changes"] = { "version": {
      "new": plugin_version[1],
      "old": plugin_version[0]
    } }

  # If not installed prepare for install.
  if plugin_state == "available":
    args.extend(["--install", name])
    comment = "Plugin installed."
    fail_comment = "Could not install plugin."
    test_comment = "Plugin will be installed."
    result["changes"] = {
      "state": {
        "new": "enabled",
        "old": "available"
      },
      "version": {
        "new": plugin_version,
        "old": None
      }
    }

  if not args:
    result["comment"] = "Unrecognised state {} for plugin.".format(
        plugin_state
    )
    result["result"]  = False
    return result

  # If __opts__["test"] return now.
  if __opts__["test"]:
    result["comment"] = test_comment
    result["result"]  = None
    return result

  # Run the command.
  full_args = [
    "--locale", "en",
    "--nogui",
    "--nosplash",
    "--modules"
  ]
  full_args.extend(args)
  output = __salt__["cmd.retcode"](bin + " " + " ".join(full_args))

def start(name, version=None, args=None):
  """Ensures that NetBeans is running.
  The command line interface is very limited and has the
  huge problem that it actually starts the IDE even when
  module commands are specified.

  If the IDE **is not** running, when the command is completed,
  the process will stay alive in full IDE mode.
  If the IDE **is** already running, when the command is completed,
  the process will terminate.

  This state should be used as a prerequisite of all other
  states so that the IDE is guaranteed to be running.

  version:
    The version of the IDE to run.

  args:
    List of arguments to pass to the command line.
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  version = _check_version(version)
  if version is None:
    result["comment"] = "Invalid version {} specified.".format(version)
    result["result"]  = False
    return result

  # Look for the plugin state.
  exceptions = __salt__["netbeans.exceptions"]()
  try:
    nb_path = __salt__["netbeans.find_installation"](version)
    bin  = path.join(nb_path, "bin", "netbeans")
    args = ["--locale", "en", "--nogui", "--nosplash"]
    code = __salt__["cmd.shell"]("nohup " + bin + " " + " ".join(args))

    if code == 0:
      result["comment"] = "Started NetBeans in silent mode."
      result["result"]  = True

    else:
      result["comment"] = "Could not start NetBeans in silent mode."
      result["result"]  = False
    return result

  except exceptions["NoInstallFound"]:
    result["comment"] = "NetBeans {} not installed.".format(version)
    result["result"]  = False
    return result

def stop(name, version=None):
  """
  """
  result = {
    "changes": None,
    "comment": "",
    "name":    name,
    "result":  False
  }

  # Resolve version.
  version = _check_version(version)
  if version is None:
    result["comment"] = "Invalid version {} specified.".format(version)
    result["result"]  = False
    return result

  # Look for the plugin state.
  exceptions   = __salt__["netbeans.exceptions"]()
  try:
    nb_path = __salt__["netbeans.find_installation"](version)
    native  = path.join(nb_path, "platform", "lib")
    cmd = (
        "ps -ef | grep '" + native +
        "' | grep -v grep | awk '{ print $2 }' | xargs kill"
    )
    __salt__["cmd.shell"](cmd)

  except exceptions["NoInstallFound"]:
    result["comment"] = "NetBeans {} not installed.".format(version)
    result["result"]  = False
    return result
<|endoftext|>


<commit_before># Copyright 2007-2016 Brian Masney <masneyb@onstation.org>
#
# 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.

# cciw/middleware/threadlocals.py
# http://code.djangoproject.com/wiki/CookBookThreadlocalsAndUser

import threading
from django.contrib.auth.models import AnonymousUser

THREAD_LOCALS = threading.local()

def get_current_user():
    return getattr(THREAD_LOCALS, 'user', None)

def get_request_uri():
    return getattr(THREAD_LOCALS, 'request_uri', None)

def get_valid_bulletins():
    return getattr(THREAD_LOCALS, 'valid_bulletins', [])

class ThreadLocals(object):
    """Middleware that gets various objects from the
    request object and saves them in thread local storage."""
    def process_request(self, request):
        THREAD_LOCALS.user = getattr(request, 'user', None)
        THREAD_LOCALS.request_uri = request.META.get('PATH_INFO')

        valid_bulletins = [1]
        user = get_current_user()

        if not isinstance(user, AnonymousUser):
            for bulletin_id in user.caveuserprofile.bulletins.get_queryset():
                valid_bulletins.append(bulletin_id.id)

        THREAD_LOCALS.valid_bulletins = valid_bulletins

<commit_msg>fixed comment style<commit_after># Copyright 2007-2016 Brian Masney <masneyb@onstation.org>
#
# 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.

# cciw/middleware/threadlocals.py
# http://code.djangoproject.com/wiki/CookBookThreadlocalsAndUser

import threading
from django.contrib.auth.models import AnonymousUser

THREAD_LOCALS = threading.local()

def get_current_user():
    return getattr(THREAD_LOCALS, 'user', None)

def get_request_uri():
    return getattr(THREAD_LOCALS, 'request_uri', None)

def get_valid_bulletins():
    return getattr(THREAD_LOCALS, 'valid_bulletins', [])

class ThreadLocals(object):
    # Middleware that gets various objects from the
    # request object and saves them in thread local storage.
    def process_request(self, request):
        THREAD_LOCALS.user = getattr(request, 'user', None)
        THREAD_LOCALS.request_uri = request.META.get('PATH_INFO')

        valid_bulletins = [1]
        user = get_current_user()

        if not isinstance(user, AnonymousUser):
            for bulletin_id in user.caveuserprofile.bulletins.get_queryset():
                valid_bulletins.append(bulletin_id.id)

        THREAD_LOCALS.valid_bulletins = valid_bulletins

<|endoftext|>


<commit_before>from __future__ import with_statement

from fabric.api import (
    shell_env,
    task,
)

def reset():
    from fabric.api import settings
    from fabtools.utils import run_as_root

    with settings(warn_only=True):
        run_as_root('apt-key del 7BD9BF62')
        run_as_root('apt-key del C4DEFFEB')

@task
def deb():
    """
    Check deb functions.
    """

    from fabric.api import run
    from fabtools import require
    from fabtools import deb

    # deb.add_apt_key with keyid
    reset()
    deb.add_apt_key(keyid='C4DEFFEB', url='http://repo.varnish-cache.org/debian/GPG-key.txt')
    run('apt-key finger | grep -q C4DEFFEB')

    reset()
    deb.add_apt_key(keyid='7BD9BF62')
    deb.add_apt_key(keyid='7BD9BF62') # Intentionally repeated
    run('apt-key finger | grep -q 7BD9BF62')

    reset()
    deb.add_apt_key(keyid='7BD9BF62', keyserver='keyserver.ubuntu.com')
    run('apt-key finger | grep -q 7BD9BF62')

    reset()
    run('wget http://repo.varnish-cache.org/debian/GPG-key.txt -O /tmp/tmp.fabtools.test.key')
    deb.add_apt_key(keyid='C4DEFFEB', filename='/tmp/tmp.fabtools.test.key')
    run('apt-key finger | grep -q C4DEFFEB')

    # deb.add_apt_key without keyid
    reset()
    deb.add_apt_key(url='http://repo.varnish-cache.org/debian/GPG-key.txt')
    run('apt-key finger | grep -q C4DEFFEB')

    reset()
    deb.add_apt_key(keyid='7BD9BF62')
    run('apt-key finger | grep -q 7BD9BF62')

    reset()
    run('wget http://repo.varnish-cache.org/debian/GPG-key.txt -O /tmp/tmp.fabtools.test.key')
    deb.add_apt_key(filename='/tmp/tmp.fabtools.test.key')
    run('apt-key finger | grep -q C4DEFFEB')

    # require.deb.key
    reset()
    require.deb.key(keyid='C4DEFFEB', url='http://repo.varnish-cache.org/debian/GPG-key.txt')
    run('apt-key finger | grep -q C4DEFFEB')

    reset()
    require.deb.key(keyid='7BD9BF62')
    require.deb.key(keyid='7BD9BF62') # Intentionally repeated
    run('apt-key finger | grep -q 7BD9BF62')

    reset()
    require.deb.key(keyid='7BD9BF62', keyserver='keyserver.ubuntu.com')
    run('apt-key finger | grep -q 7BD9BF62')

    reset()
    run('wget http://repo.varnish-cache.org/debian/GPG-key.txt -O /tmp/tmp.fabtools.test.key')
    require.deb.key(keyid='C4DEFFEB', filename='/tmp/tmp.fabtools.test.key')
    run('apt-key finger | grep -q C4DEFFEB')

<commit_msg>Fix test fabfile to work on Debian squeeze [#141]<commit_after>from __future__ import with_statement

from fabric.api import (
    run,
    task,
)

def reset():
    from fabric.api import settings
    from fabtools.utils import run_as_root

    with settings(warn_only=True):
        run_as_root('apt-key del 7BD9BF62')
        run_as_root('apt-key del C4DEFFEB')

@task
def deb():
    """
    Check deb functions.
    """

    from fabtools import require
    from fabtools import deb
    from fabtools.utils import run_as_root

    # deb.add_apt_key with keyid
    reset()
    deb.add_apt_key(keyid='C4DEFFEB', url='http://repo.varnish-cache.org/debian/GPG-key.txt')
    run_as_root('apt-key finger | grep -q C4DEFFEB')

    reset()
    deb.add_apt_key(keyid='7BD9BF62')
    deb.add_apt_key(keyid='7BD9BF62') # Intentionally repeated
    run_as_root('apt-key finger | grep -q 7BD9BF62')

    reset()
    deb.add_apt_key(keyid='7BD9BF62', keyserver='keyserver.ubuntu.com')
    run_as_root('apt-key finger | grep -q 7BD9BF62')

    reset()
    run('wget http://repo.varnish-cache.org/debian/GPG-key.txt -O /tmp/tmp.fabtools.test.key')
    deb.add_apt_key(keyid='C4DEFFEB', filename='/tmp/tmp.fabtools.test.key')
    run_as_root('apt-key finger | grep -q C4DEFFEB')

    # deb.add_apt_key without keyid
    reset()
    deb.add_apt_key(url='http://repo.varnish-cache.org/debian/GPG-key.txt')
    run_as_root('apt-key finger | grep -q C4DEFFEB')

    reset()
    deb.add_apt_key(keyid='7BD9BF62')
    run_as_root('apt-key finger | grep -q 7BD9BF62')

    reset()
    run('wget http://repo.varnish-cache.org/debian/GPG-key.txt -O /tmp/tmp.fabtools.test.key')
    deb.add_apt_key(filename='/tmp/tmp.fabtools.test.key')
    run_as_root('apt-key finger | grep -q C4DEFFEB')

    # require.deb.key
    reset()
    require.deb.key(keyid='C4DEFFEB', url='http://repo.varnish-cache.org/debian/GPG-key.txt')
    run_as_root('apt-key finger | grep -q C4DEFFEB')

    reset()
    require.deb.key(keyid='7BD9BF62')
    require.deb.key(keyid='7BD9BF62') # Intentionally repeated
    run_as_root('apt-key finger | grep -q 7BD9BF62')

    reset()
    require.deb.key(keyid='7BD9BF62', keyserver='keyserver.ubuntu.com')
    run_as_root('apt-key finger | grep -q 7BD9BF62')

    reset()
    run('wget http://repo.varnish-cache.org/debian/GPG-key.txt -O /tmp/tmp.fabtools.test.key')
    require.deb.key(keyid='C4DEFFEB', filename='/tmp/tmp.fabtools.test.key')
    run_as_root('apt-key finger | grep -q C4DEFFEB')

<|endoftext|>