<commit_before>import logging
import os
import sys
from pathlib import Path
from subprocess import check_output, CalledProcessError

import numpy as np
import pandas as pd

from emgapianns.management.lib import utils
from emgapi import models as emg_models
from django.db.models import Q

class StudySummaryGenerator(object):

    def __init__(self, accession, pipeline, rootpath, database):
        self.study_accession = accession
        self.pipeline = pipeline
        self.rootpath = rootpath
        self.emg_db_name = database
        self.study = emg_models.Study.objects.using(self.emg_db_name).get(secondary_accession=self.study_accession)
        self.study_result_dir = os.path.join(self.rootpath, self.study.result_directory)
        self.summary_dir = None

    def run(self):
        if not os.path.exists(self.study_result_dir):
            sys.exit(
                f"Study result directory for {self.study_accession} does not exist:\n{self.study_result_dir}")

        jobs = emg_models.AnalysisJob.objects.using(self.emg_db_name)
        jobs = jobs.filter(
            Q(study__secondary_accession=self.study_accession) &
            Q(analysis_status__analysis_status='completed') &
            Q(pipeline__release_version=self.pipeline))

        analysis_result_dirs = {}
        for job in jobs:
            job_result_directory = os.path.join(self.rootpath, job.result_directory)
            job_input_file_name = job.input_file_name
            analysis_result_dirs[job_input_file_name] = job_result_directory

        self.summary_dir = os.path.join(self.study_result_dir, 'version_{}/project-summary'.format(self.pipeline))
        self.create_summary_dir()

        # TODO: Add UNITE and ITSoneDB
        if self.pipeline == '4.1':
            self.generate_taxonomy_phylum_summary(analysis_result_dirs, self.pipeline, 'SSU',
                                                  'phylum_taxonomy_abundances_SSU_v{}.tsv'.format(self.pipeline))
            self.generate_taxonomy_phylum_summary(analysis_result_dirs, self.pipeline, 'LSU',
                                                  f'phylum_taxonomy_abundances_LSU_v{self.pipeline}.tsv')

        self.generate_taxonomy_summary(analysis_result_dirs, 'SSU',
                                       'taxonomy_abundances_SSU_v{}.tsv'.format(self.pipeline))
        self.generate_taxonomy_summary(analysis_result_dirs, 'LSU',
                                       'taxonomy_abundances_LSU_v{}.tsv'.format(self.pipeline))

        self.generate_ipr_summary(analysis_result_dirs, 'IPR_abundances_v{}.tsv'.format(self.pipeline))

        self.generate_go_summary(analysis_result_dirs, 'slim')
        self.generate_go_summary(analysis_result_dirs, 'full')

        logging.info("Program finished successfully.")

    def generate_taxonomy_phylum_summary(self, analysis_result_dirs, version, su_type, filename):
        if version == '4.1':
            self.generate_taxonomy_phylum_summary_v4(su_type, filename)
        else:
            self.generate_taxonomy_phylum_summary_v5(analysis_result_dirs, su_type, filename)

    def generate_taxonomy_phylum_summary_v4(self, su_type, filename):
        res_file_re = os.path.join('**', 'taxonomy-summary', su_type, 'kingdom-counts.txt')
        res_files = self.get_raw_result_files(res_file_re)
        study_df = self.merge_dfs(res_files,
                                  delimiter='\t',
                                  key=['kingdom', 'phylum'],
                                  raw_cols=['kingdom', 'phylum', 'count', 'ignored'])
        self.write_results_file(study_df, filename)

        alias = '{}_phylum_taxonomy_abundances_{}_v{}.tsv'.format(self.study_accession, su_type, self.pipeline)
        description = 'Phylum level taxonomies {}'.format(su_type)
        group = 'Taxonomic analysis {} rRNA'.format(su_type)
        self.upload_study_file(filename, alias, description, group)

    def generate_taxonomy_phylum_summary_v5(self, analysis_result_dirs, su_type, filename):
        # Fixme: Work in progress! Finish implementation
        res_files = self.get_mapseq_result_files(analysis_result_dirs, su_type)

        study_df = self.merge_dfs(res_files,
                                  delimiter='\t',
                                  key=['kingdom', 'phylum'],
                                  raw_cols=['kingdom', 'phylum', 'count', 'ignored'])
        self.write_results_file(study_df, filename)

        alias = '{}_phylum_taxonomy_abundances_{}_v{}.tsv'.format(self.study_accession, su_type, self.pipeline)
        description = 'Phylum level taxonomies {}'.format(su_type)
        group = 'Taxonomic analysis {} rRNA'.format(su_type)
        self.upload_study_file(filename, alias, description, group)

    def generate_taxonomy_summary(self, analysis_result_dirs, su_type, filename):
        res_files = self.get_taxonomy_result_files(analysis_result_dirs, su_type)
        study_df = self.merge_dfs(res_files,
                                  key=['lineage'],
                                  delimiter='\t',
                                  raw_cols=['OTU', 'count', 'lineage'],
                                  skip_rows=2)
        study_df = study_df.rename(columns={'lineage': '#SampleID'})
        self.write_results_file(study_df, filename)

        alias = '{}_taxonomy_abundances_{}_v{}.tsv'.format(self.study_accession, su_type, self.pipeline)
        description = 'Taxonomic assignments {}'.format(su_type)
        group = 'Taxonomic analysis {} rRNA'.format(su_type)
        self.upload_study_file(filename, alias, description, group)

    def get_raw_result_files(self, res_file_re):
        paths = list(Path(self.study_result_dir).glob(res_file_re))
        return [str(p.resolve()) for p in paths]

    def merge_dfs(self, filelist, delimiter, key, raw_cols, skip_rows=0):
        study_df = pd.DataFrame(columns=key)
        for f in sorted(filelist):
            accession = utils.get_accession_from_result_dir_path(f)
            df = self.read_count_tsv(f, delimiter, raw_cols, skip_rows)
            df = df.filter(key + ['count'])
            df = df.rename(columns={'count': accession})
            study_df = study_df.merge(df, on=key, how='outer')
        study_df = study_df.sort_values(by=key)
        study_df = self.clean_summary_df(study_df)
        return study_df

    def generate_ipr_summary(self, analysis_result_dirs, filename):
        res_files = self.get_ipr_result_files(analysis_result_dirs)
        study_df = self.merge_dfs(res_files, delimiter=',',
                                  key=['IPR', 'description'],
                                  raw_cols=['IPR', 'description', 'count'])
        self.write_results_file(study_df, filename)

        alias = '{}_IPR_abundances_v{}.tsv'.format(self.study_accession, self.pipeline)
        description = 'InterPro matches'
        self.upload_study_file(filename, alias, description, 'Functional analysis')

    def generate_go_summary(self, analysis_result_dirs, mode):
        if mode == 'slim':
            sum_file = 'GO-slim'
            description = 'GO slim annotation'
        else:
            sum_file = 'GO'
            description = 'Complete GO annotation'

        res_files = self.get_go_result_files(analysis_result_dirs, mode)
        study_df = self.merge_dfs(res_files, delimiter=',',
                                  key=['GO', 'description', 'category'],
                                  raw_cols=['GO', 'description', 'category', 'count'])
        study_df['description'] = study_df['description'].str.replace(',', '@')
        study_df['category'] = study_df['category'].str.replace('_', ' ')
        realname = sum_file + '_abundances_v{}.tsv'.format(self.pipeline)
        self.write_results_file(study_df, realname)

        self.generate_filtered_go_summary(study_df,
                                          'category == "cellular component"',
                                          'CC_{}_abundances_v{}.tsv'.format(sum_file, self.pipeline))

        self.generate_filtered_go_summary(study_df,
                                          'category == "biological process"',
                                          'BP_{}_abundances_v{}.tsv'.format(sum_file, self.pipeline))

        self.generate_filtered_go_summary(study_df,
                                          'category not in ["biological process", "cellular component"]',
                                          'MF_{}_abundances_v{}.tsv'.format(sum_file, self.pipeline))

        alias = '{}_{}_abundances_v{}.tsv'.format(self.study_accession, sum_file, self.pipeline)
        self.upload_study_file(realname, alias, description, 'Functional analysis')
        return study_df

    def generate_filtered_go_summary(self, df, query, filename):
        df = df.query(query)
        self.write_results_file(df, filename)

    @staticmethod
    def read_count_tsv(filename, delimiter, cols, skip_rows=0):
        df = pd.read_csv(filename, delimiter=delimiter, names=cols, skiprows=skip_rows)
        df = df.astype({'count': 'int'})
        return df

    def write_results_file(self, df, filename):
        filepath = os.path.join(self.summary_dir, filename)
        df.to_csv(filepath, sep='\t', header=True, index=False)

    @staticmethod
    def clean_summary_df(df):
        df = df.fillna(0)
        int_cols = df.select_dtypes(include=[np.number]).columns.tolist()
        df = df.astype({col: 'int' for col in int_cols})
        return df

    def create_summary_dir(self):
        os.makedirs(self.summary_dir, exist_ok=True)

    def upload_study_file(self, realname, alias, description, group):
        file_config = {
            'alias': alias,
            'compression': False,
            'description_label': description,
            'format_name': 'TSV',
            'group_type': group,
            'real_name': realname,
            'subdir': 'version_{}/project-summary'.format(self.pipeline),
            '_required': True
        }
        _study_rootpath = self.study_result_dir.replace('version_{}'.format(self.pipeline), '')
        utils.StudyDownload(self.emg_db_name, _study_rootpath, file_config, self.pipeline).save(self.study)

    @staticmethod
    def get_mapseq_result_files(analysis_result_dirs, su_type):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            res_file_re = os.path.join(dir, 'taxonomy-summary', su_type,
                                       '{}_{}.fasta.mseq.gz'.format(input_file_name, su_type))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def get_taxonomy_result_files(analysis_result_dirs, su_type):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            res_file_re = os.path.join(dir, 'taxonomy-summary', su_type,
                                       '{}_{}.fasta.mseq.tsv'.format(input_file_name, su_type))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def get_ipr_result_files(analysis_result_dirs):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            res_file_re = os.path.join(dir, '{}_summary.ipr'.format(input_file_name))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def get_go_result_files(analysis_result_dirs, mode):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            file_name = '{}_summary.go' if mode == 'full' else '{}_summary.go_slim'
            res_file_re = os.path.join(dir, file_name.format(input_file_name))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def __get_phylum_counts(filepath):
        """
            Summarises the counts for each unique phylum.
        :return:
        """
        std_out = check_output(
            "zcat ERR3506537_MERGED_FASTQ_SSU.fasta.mseq.gz | grep -v '^#' | cut -f 14- | cut -d ';' -f 1-3 | sed 's/\t$//' | sed 's/;p__$//' | sed 's/;k__$//' | sort | uniq -c".format(
                filepath), shell=True).rstrip()
        return std_out

    @staticmethod
    def __parse_phylum_counts(std_out):
        """
            Parse the standard output of the Linux oneliner.
        :return:
        """
        # TODO: Implement
        pass

    @staticmethod
    def __count_number_of_lines(filepath):
        """
            Counts number of lines in compressed text file.
        :return:
        """
        count = check_output("zcat {} | wc -l".format(filepath), shell=True).rstrip()
        return int(count)

    @staticmethod
    def __count_number_of_seqs(filepath):
        """
            Counts number of sequences in compressed fasta file.
        :return:
        """
        try:
            count = check_output("zcat {} | grep -c '>'".format(filepath), shell=True).rstrip()
            return int(count)
        except CalledProcessError:
            return 0
<commit_msg>Added exception for amplicon datasets.<commit_after>import logging
import os
import sys
from pathlib import Path
from subprocess import check_output, CalledProcessError

import numpy as np
import pandas as pd

from emgapianns.management.lib import utils
from emgapi import models as emg_models
from django.db.models import Q

class StudySummaryGenerator(object):

    def __init__(self, accession, pipeline, rootpath, database):
        self.study_accession = accession
        self.pipeline = pipeline
        self.rootpath = rootpath
        self.emg_db_name = database
        self.study = emg_models.Study.objects.using(self.emg_db_name).get(secondary_accession=self.study_accession)
        self.study_result_dir = os.path.join(self.rootpath, self.study.result_directory)
        self.summary_dir = None

    def run(self):
        if not os.path.exists(self.study_result_dir):
            sys.exit(
                f"Study result directory for {self.study_accession} does not exist:\n{self.study_result_dir}")

        jobs = emg_models.AnalysisJob.objects.using(self.emg_db_name)
        jobs = jobs.filter(
            Q(study__secondary_accession=self.study_accession) &
            Q(analysis_status__analysis_status='completed') &
            Q(pipeline__release_version=self.pipeline))

        experiment_types = set()
        analysis_result_dirs = {}
        for job in jobs:
            job_result_directory = os.path.join(self.rootpath, job.result_directory)
            job_input_file_name = job.input_file_name
            analysis_result_dirs[job_input_file_name] = job_result_directory
            experiment_types.add(job.experiment_type.experiment_type)

        self.summary_dir = os.path.join(self.study_result_dir, 'version_{}/project-summary'.format(self.pipeline))
        self.create_summary_dir()

        # TODO: Add UNITE and ITSoneDB
        if self.pipeline == '4.1':
            self.generate_taxonomy_phylum_summary(analysis_result_dirs, self.pipeline, 'SSU',
                                                  'phylum_taxonomy_abundances_SSU_v{}.tsv'.format(self.pipeline))
            self.generate_taxonomy_phylum_summary(analysis_result_dirs, self.pipeline, 'LSU',
                                                  f'phylum_taxonomy_abundances_LSU_v{self.pipeline}.tsv')

        self.generate_taxonomy_summary(analysis_result_dirs, 'SSU',
                                       'taxonomy_abundances_SSU_v{}.tsv'.format(self.pipeline))
        self.generate_taxonomy_summary(analysis_result_dirs, 'LSU',
                                       'taxonomy_abundances_LSU_v{}.tsv'.format(self.pipeline))

        if len(experiment_types) == 1 and 'amplicon' in experiment_types:
            logging.info("AMPLICON datasets only! Skipping the generation of the functional matrix files!")
        else:
            self.generate_ipr_summary(analysis_result_dirs, 'IPR_abundances_v{}.tsv'.format(self.pipeline))
            self.generate_go_summary(analysis_result_dirs, 'slim')
            self.generate_go_summary(analysis_result_dirs, 'full')

        logging.info("Program finished successfully.")

    def generate_taxonomy_phylum_summary(self, analysis_result_dirs, version, su_type, filename):
        if version == '4.1':
            self.generate_taxonomy_phylum_summary_v4(su_type, filename)
        else:
            self.generate_taxonomy_phylum_summary_v5(analysis_result_dirs, su_type, filename)

    def generate_taxonomy_phylum_summary_v4(self, su_type, filename):
        res_file_re = os.path.join('**', 'taxonomy-summary', su_type, 'kingdom-counts.txt')
        res_files = self.get_raw_result_files(res_file_re)
        study_df = self.merge_dfs(res_files,
                                  delimiter='\t',
                                  key=['kingdom', 'phylum'],
                                  raw_cols=['kingdom', 'phylum', 'count', 'ignored'])
        self.write_results_file(study_df, filename)

        alias = '{}_phylum_taxonomy_abundances_{}_v{}.tsv'.format(self.study_accession, su_type, self.pipeline)
        description = 'Phylum level taxonomies {}'.format(su_type)
        group = 'Taxonomic analysis {} rRNA'.format(su_type)
        self.upload_study_file(filename, alias, description, group)

    def generate_taxonomy_phylum_summary_v5(self, analysis_result_dirs, su_type, filename):
        # Fixme: Work in progress! Finish implementation
        res_files = self.get_mapseq_result_files(analysis_result_dirs, su_type)

        study_df = self.merge_dfs(res_files,
                                  delimiter='\t',
                                  key=['kingdom', 'phylum'],
                                  raw_cols=['kingdom', 'phylum', 'count', 'ignored'])
        self.write_results_file(study_df, filename)

        alias = '{}_phylum_taxonomy_abundances_{}_v{}.tsv'.format(self.study_accession, su_type, self.pipeline)
        description = 'Phylum level taxonomies {}'.format(su_type)
        group = 'Taxonomic analysis {} rRNA'.format(su_type)
        self.upload_study_file(filename, alias, description, group)

    def generate_taxonomy_summary(self, analysis_result_dirs, su_type, filename):
        res_files = self.get_taxonomy_result_files(analysis_result_dirs, su_type)
        study_df = self.merge_dfs(res_files,
                                  key=['lineage'],
                                  delimiter='\t',
                                  raw_cols=['OTU', 'count', 'lineage'],
                                  skip_rows=2)
        study_df = study_df.rename(columns={'lineage': '#SampleID'})
        self.write_results_file(study_df, filename)

        alias = '{}_taxonomy_abundances_{}_v{}.tsv'.format(self.study_accession, su_type, self.pipeline)
        description = 'Taxonomic assignments {}'.format(su_type)
        group = 'Taxonomic analysis {} rRNA'.format(su_type)
        self.upload_study_file(filename, alias, description, group)

    def get_raw_result_files(self, res_file_re):
        paths = list(Path(self.study_result_dir).glob(res_file_re))
        return [str(p.resolve()) for p in paths]

    def merge_dfs(self, filelist, delimiter, key, raw_cols, skip_rows=0):
        study_df = pd.DataFrame(columns=key)
        for f in sorted(filelist):
            accession = utils.get_accession_from_result_dir_path(f)
            df = self.read_count_tsv(f, delimiter, raw_cols, skip_rows)
            df = df.filter(key + ['count'])
            df = df.rename(columns={'count': accession})
            study_df = study_df.merge(df, on=key, how='outer')
        study_df = study_df.sort_values(by=key)
        study_df = self.clean_summary_df(study_df)
        return study_df

    def generate_ipr_summary(self, analysis_result_dirs, filename):
        res_files = self.get_ipr_result_files(analysis_result_dirs)
        study_df = self.merge_dfs(res_files, delimiter=',',
                                  key=['IPR', 'description'],
                                  raw_cols=['IPR', 'description', 'count'])
        self.write_results_file(study_df, filename)

        alias = '{}_IPR_abundances_v{}.tsv'.format(self.study_accession, self.pipeline)
        description = 'InterPro matches'
        self.upload_study_file(filename, alias, description, 'Functional analysis')

    def generate_go_summary(self, analysis_result_dirs, mode):
        if mode == 'slim':
            sum_file = 'GO-slim'
            description = 'GO slim annotation'
        else:
            sum_file = 'GO'
            description = 'Complete GO annotation'

        res_files = self.get_go_result_files(analysis_result_dirs, mode)
        study_df = self.merge_dfs(res_files, delimiter=',',
                                  key=['GO', 'description', 'category'],
                                  raw_cols=['GO', 'description', 'category', 'count'])
        study_df['description'] = study_df['description'].str.replace(',', '@')
        study_df['category'] = study_df['category'].str.replace('_', ' ')
        realname = sum_file + '_abundances_v{}.tsv'.format(self.pipeline)
        self.write_results_file(study_df, realname)

        self.generate_filtered_go_summary(study_df,
                                          'category == "cellular component"',
                                          'CC_{}_abundances_v{}.tsv'.format(sum_file, self.pipeline))

        self.generate_filtered_go_summary(study_df,
                                          'category == "biological process"',
                                          'BP_{}_abundances_v{}.tsv'.format(sum_file, self.pipeline))

        self.generate_filtered_go_summary(study_df,
                                          'category not in ["biological process", "cellular component"]',
                                          'MF_{}_abundances_v{}.tsv'.format(sum_file, self.pipeline))

        alias = '{}_{}_abundances_v{}.tsv'.format(self.study_accession, sum_file, self.pipeline)
        self.upload_study_file(realname, alias, description, 'Functional analysis')
        return study_df

    def generate_filtered_go_summary(self, df, query, filename):
        df = df.query(query)
        self.write_results_file(df, filename)

    @staticmethod
    def read_count_tsv(filename, delimiter, cols, skip_rows=0):
        df = pd.read_csv(filename, delimiter=delimiter, names=cols, skiprows=skip_rows)
        df = df.astype({'count': 'int'})
        return df

    def write_results_file(self, df, filename):
        filepath = os.path.join(self.summary_dir, filename)
        df.to_csv(filepath, sep='\t', header=True, index=False)

    @staticmethod
    def clean_summary_df(df):
        df = df.fillna(0)
        int_cols = df.select_dtypes(include=[np.number]).columns.tolist()
        df = df.astype({col: 'int' for col in int_cols})
        return df

    def create_summary_dir(self):
        os.makedirs(self.summary_dir, exist_ok=True)

    def upload_study_file(self, realname, alias, description, group):
        file_config = {
            'alias': alias,
            'compression': False,
            'description_label': description,
            'format_name': 'TSV',
            'group_type': group,
            'real_name': realname,
            'subdir': 'version_{}/project-summary'.format(self.pipeline),
            '_required': True
        }
        _study_rootpath = self.study_result_dir.replace('version_{}'.format(self.pipeline), '')
        utils.StudyDownload(self.emg_db_name, _study_rootpath, file_config, self.pipeline).save(self.study)

    @staticmethod
    def get_mapseq_result_files(analysis_result_dirs, su_type):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            res_file_re = os.path.join(dir, 'taxonomy-summary', su_type,
                                       '{}_{}.fasta.mseq.gz'.format(input_file_name, su_type))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def get_taxonomy_result_files(analysis_result_dirs, su_type):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            res_file_re = os.path.join(dir, 'taxonomy-summary', su_type,
                                       '{}_{}.fasta.mseq.tsv'.format(input_file_name, su_type))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def get_ipr_result_files(analysis_result_dirs):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            res_file_re = os.path.join(dir, '{}_summary.ipr'.format(input_file_name))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def get_go_result_files(analysis_result_dirs, mode):
        result = []
        for input_file_name, dir in analysis_result_dirs.items():
            file_name = '{}_summary.go' if mode == 'full' else '{}_summary.go_slim'
            res_file_re = os.path.join(dir, file_name.format(input_file_name))
            if os.path.exists(res_file_re):
                result.append(res_file_re)
            else:
                logging.warning("Result file does not exist:\n{}".format(res_file_re))
        return result

    @staticmethod
    def __get_phylum_counts(filepath):
        """
            Summarises the counts for each unique phylum.
        :return:
        """
        std_out = check_output(
            "zcat ERR3506537_MERGED_FASTQ_SSU.fasta.mseq.gz | grep -v '^#' | cut -f 14- | cut -d ';' -f 1-3 | sed 's/\t$//' | sed 's/;p__$//' | sed 's/;k__$//' | sort | uniq -c".format(
                filepath), shell=True).rstrip()
        return std_out

    @staticmethod
    def __parse_phylum_counts(std_out):
        """
            Parse the standard output of the Linux oneliner.
        :return:
        """
        # TODO: Implement
        pass

    @staticmethod
    def __count_number_of_lines(filepath):
        """
            Counts number of lines in compressed text file.
        :return:
        """
        count = check_output("zcat {} | wc -l".format(filepath), shell=True).rstrip()
        return int(count)

    @staticmethod
    def __count_number_of_seqs(filepath):
        """
            Counts number of sequences in compressed fasta file.
        :return:
        """
        try:
            count = check_output("zcat {} | grep -c '>'".format(filepath), shell=True).rstrip()
            return int(count)
        except CalledProcessError:
            return 0
<|endoftext|>


<commit_before># mqttclient application command

from django.core.management.base import BaseCommand
from django.conf import settings

from ambulances.management.commands._client import BaseClient

from ambulances.models import EquipmentCount, Ambulances, Status, Call
from ambulances.serializers import MQTTLocationSerializer, MQTTAmbulanceLocSerializer, CallSerializer

from django.utils.six import BytesIO
from django.core.exceptions import ObjectDoesNotExist
from django.contrib.gis.geos import Point
from rest_framework.parsers import JSONParser
from rest_framework.renderers import JSONRenderer

# Client
class Client(BaseClient):

    # The callback for when the client receives a CONNACK
    # response from the server.
    def on_connect(self, client, userdata, flags, rc):

        # is connected?
        if not super().on_connect(client, userdata, flags, rc):
            return False

        # Subscribing in on_connect() means that if we lose the
        # connection and reconnect then subscriptions will be renewed.
        client.subscribe('#', 2)

        # hospital message handler
        self.client.message_callback_add('hospital/+/equipment/#',
                                         self.on_hospital)

        self.client.message_callback_add('ambulance/+/call', self.on_call)

        # user location handler
        self.client.message_callback_add('user/+/location', self.on_user_loc)

        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(">> Listening to messages..."))

        return True

    # The callback for when a PUBLISH message is received from the server.
    def on_message(self, client, userdata, msg):

        # handle unknown messages
        self.stdout.write(
            self.style.WARNING(
                "*> Unknown message topic {} {}".format(msg.topic,
                                                        msg.payload)))

    # Update hospital resources
    def on_hospital(self, client, userdata, msg):

        # parse message
        topic = msg.topic.split('/')
        hospital = topic[1]
        equipment = topic[3]

        if not msg.payload:
            return

        try:
            quantity = int(msg.payload)
            if self.verbosity > 1:
                self.stdout.write(" > {} {}".format(msg.topic, quantity))

            e = EquipmentCount.objects.get(hospital=hospital,
                                           equipment__name=equipment)

            # update quantity
            if e.quantity != quantity:
                e.quantity = quantity
                e.save()
                if self.verbosity > 0:
                    self.stdout.write(
                        self.style.SUCCESS(">> Hospital '{}' equipment '{}' updated to '{}'".format(e.hospital.name, equipment, quantity)))

        except Exception:
            self.stdout.write(
                self.style.ERROR("*> hospital/{}/equipment/{} is not available".format(hospital, equipment)))

    # Update call on database or create new call if doesn't exist
    def on_call(self, client, userdata, msg):
        # parse message
        topic = msg.topic.split('/')
        ambulance_id = topic[1]

        if not msg.payload:
            return

        try:
            # Parse data into json dict
            stream = BytesIO(msg.payload)
            data = JSONParser().parse(stream)
        except Exception:
            self.stdout.write(self.style.ERROR("ERROR PARSING CALL JSON"))

        # Try to update existing call if it exists
        try:
            name = data['name']
            residential_unit = data['residential_unit']
            stmain_number = data['stmain_number']
            delegation = data['delegation']
            zipcode = data['zipcode']
            city = data['city']
            state = data['state']
            latitude = data['latitude']
            longitude = data['longitude']
            assignment = data['assignment']
            description = data['description']
            call_time = data['call_time']
            departure_time = data['departure_time']
            transfer_time = data['transfer_time']
            hospital_time = data['hospital_time']
            base_time = data['base_time']
            ambulance = data['ambulance']

            # Obtain call object
            call = Call.objects.filter(ambulance=ambulance).first()

            # update name
            if call.name != name:
                call.name = name

            # update residential_unit
            if call.residential_unit != residential_unit:
                call.residential_unit = residential_unit

            # update stmain_number
            if call.stmain_number != stmain_number:
                call.stmain_number = stmain_number

            # update delegation
            if call.delegation != delegation:
                call.delegation = delegation

            # update zipcode
            if call.zipcode != zipcode:
                call.zipcode = zipcode

            # update city
            if call.city != city:
                call.city = city

            # update state
            if call.state != state:
                call.state = state

            # update location
            if call.location.x != longitude or call.location.y != latitude:
                call.location = Point(longitude, latitude, srid=4326)

            # update assignment
            if call.assignment != assignment:
                call.assignment = assignment

            # update description
            if call.description != description:
                call.description = description

            # update call_time
            if call.call_time != call_time:
                call.call_time = call_time

            # update departure_time
            if call.departure_time != departure_time:
                call.departure_time = departure_time

            # update transfer_time
            if call.transfer_time != transfer_time:
                call.transfer_time = transfer_time

            # update hospital_time
            if call.hospital_time != hospital_time:
                call.hospital_time = hospital_time

            # update base_time
            if call.base_time != base_time:
                call.base_time = base_time

            call.save()
            self.stdout.write(self.style.SUCCESS("Updated Call to database"))

        # If object does not exist on database, create it
        except ObjectDoesNotExist:
            serializer = CallSerializer(data=data)

            if serializer.is_valid():
               try:
                   call = serializer.save()
                   self.stdout.write(self.style.SUCCESS("Wrote Call to database"))
               except Exception:
                   self.stdout.write(self.style.ERROR("Error writing Call to database"))
            else:
                self.stdout.write("Serializer Not Valid" + str(data))

        except Exception as e:
            self.stdout.write(self.style.ERROR("Error with parsing call fields"))

    # Update user location
    def on_user_loc(self, client, userdata, msg):

        topic = msg.topic.split('/')
        user = topic[1]

        if not msg.payload:
            return

        # Parse data into json dict
        stream = BytesIO(msg.payload)
        data = JSONParser().parse(stream)

        # TODO Find out which ambulance is linked to user
        ambulance = 1
        data['ambulance'] = ambulance

        # Serialize data into object
        serializer = MQTTLocationSerializer(data=data)

        if serializer.is_valid():
            try:
                # MAURICIO: Does that save the location to the database? Are we saving the user location as well?
                lp = serializer.save()

                # Publish ambulance location as soon as user location saved
                try:
                    self.pub_ambulance_loc(client, ambulance, lp)
                except Exception as e:
                    self.stdout.write(
                        self.style.ERROR("*> Failed to publish location. Exception: {}".format(e)))

                self.stdout.write(
                        self.style.SUCCESS(">> LocationPoint for user {} in ambulance {} successfully created.".format(user, ambulance)))
            except Exception as e:
                self.stdout.write(
                    self.style.ERROR("*> LocationPoint for user {} in ambulance {} failed to create. Exception: {}".format(user, ambulance, e)))

        else:
            self.stdout.write(
                self.style.ERROR("*> Input data for user location invalid"))

    # Publish location for ambulance - after ambulance starts querying location table,
    # take out the dependency on the lp generated
    def pub_ambulance_loc(self, client, amb_id, lp):
        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(">> Publishing location for ambulance {}".format(amb_id)))

        try:
            ambulance = Ambulances.objects.get(id=amb_id)

            # Set status and grab latest location; status should be calculated
            # rather than what this code is doing
            # status = Status.objects.all().first()
            # ambulance.status = status

            # Calculate ambulance orientation?
            ambulance.orientation = 200

            # Save changes made to ambulance
            ambulance.save()

            # Convert obj back to json
            serializer = MQTTAmbulanceLocSerializer(ambulance)
            json = JSONRenderer().render(serializer.data)

            # Publish json - be sure to do this in the seeder
            client.publish('ambulance/{}/location'.format(amb_id), json, qos=2, retain=True)
            client.publish('ambulance/{}/status'.format(amb_id), ambulance.status.name, qos=2, retain=True)

        except Exception as e:
            print(e)
            self.stdout.write(
                self.style.ERROR("*> Ambulance {} does not exist".format(amb_id)))

class Command(BaseCommand):
    help = 'Connect to the mqtt broker'

    def handle(self, *args, **options):

        broker = {
            'USERNAME': '',
            'PASSWORD': '',
            'HOST': 'localhost',
            'PORT': 1883,
            'KEEPALIVE': 60,
            'CLIENT_ID': 'django',
            'CLEAN_SESSION': True
        }
        broker.update(settings.MQTT)

        client = Client(broker, self.stdout, self.style,
                        verbosity = options['verbosity'])

        self.stdout.write(
            self.style.SUCCESS("""* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
* * *                    M Q T T   C L I E N T                    * * *
* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"""))

        try:
            client.loop_forever()

        except KeyboardInterrupt:
            pass

        finally:
            client.disconnect()
<commit_msg>Implemented MQTT status publish<commit_after># mqttclient application command

from django.core.management.base import BaseCommand
from django.conf import settings

from ambulances.management.commands._client import BaseClient

from ambulances.models import EquipmentCount, Ambulances, Status, Call
from ambulances.serializers import MQTTLocationSerializer, MQTTAmbulanceLocSerializer, CallSerializer

from django.utils.six import BytesIO
from django.core.exceptions import ObjectDoesNotExist
from django.contrib.gis.geos import Point
from rest_framework.parsers import JSONParser
from rest_framework.renderers import JSONRenderer

# Client
class Client(BaseClient):

    # The callback for when the client receives a CONNACK
    # response from the server.
    def on_connect(self, client, userdata, flags, rc):

        # is connected?
        if not super().on_connect(client, userdata, flags, rc):
            return False

        # Subscribing in on_connect() means that if we lose the
        # connection and reconnect then subscriptions will be renewed.
        client.subscribe('#', 2)

        # hospital message handler
        self.client.message_callback_add('hospital/+/equipment/#',
                                         self.on_hospital)

        # call handler
        self.client.message_callback_add('ambulance/+/call', self.on_call)

        # user location handler
        self.client.message_callback_add('user/+/location', self.on_user_loc)

        # status handler
        self.client.message_callback_add('ambulance/+/status', self.on_status)

        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(">> Listening to messages..."))

        return True

    # The callback for when a PUBLISH message is received from the server.
    def on_message(self, client, userdata, msg):

        # handle unknown messages
        self.stdout.write(
            self.style.WARNING(
                "*> Unknown message topic {} {}".format(msg.topic,
                                                        msg.payload)))

    # Update hospital resources
    def on_hospital(self, client, userdata, msg):

        # parse message
        topic = msg.topic.split('/')
        hospital = topic[1]
        equipment = topic[3]

        if not msg.payload:
            return

        try:
            quantity = int(msg.payload)
            if self.verbosity > 1:
                self.stdout.write(" > {} {}".format(msg.topic, quantity))

            e = EquipmentCount.objects.get(hospital=hospital,
                                           equipment__name=equipment)

            # update quantity
            if e.quantity != quantity:
                e.quantity = quantity
                e.save()
                if self.verbosity > 0:
                    self.stdout.write(
                        self.style.SUCCESS(">> Hospital '{}' equipment '{}' updated to '{}'".format(e.hospital.name, equipment, quantity)))

        except Exception:
            self.stdout.write(
                self.style.ERROR("*> hospital/{}/equipment/{} is not available".format(hospital, equipment)))

    # Update call on database or create new call if doesn't exist
    def on_call(self, client, userdata, msg):
        # parse message
        topic = msg.topic.split('/')
        ambulance_id = topic[1]

        if not msg.payload:
            return

        try:
            # Parse data into json dict
            stream = BytesIO(msg.payload)
            data = JSONParser().parse(stream)
        except Exception:
            self.stdout.write(self.style.ERROR("ERROR PARSING CALL JSON"))

        # Try to update existing call if it exists
        try:
            name = data['name']
            residential_unit = data['residential_unit']
            stmain_number = data['stmain_number']
            delegation = data['delegation']
            zipcode = data['zipcode']
            city = data['city']
            state = data['state']
            latitude = data['latitude']
            longitude = data['longitude']
            assignment = data['assignment']
            description = data['description']
            call_time = data['call_time']
            departure_time = data['departure_time']
            transfer_time = data['transfer_time']
            hospital_time = data['hospital_time']
            base_time = data['base_time']
            ambulance = data['ambulance']

            # Obtain call object
            call = Call.objects.filter(ambulance=ambulance).first()

            # update name
            if call.name != name:
                call.name = name

            # update residential_unit
            if call.residential_unit != residential_unit:
                call.residential_unit = residential_unit

            # update stmain_number
            if call.stmain_number != stmain_number:
                call.stmain_number = stmain_number

            # update delegation
            if call.delegation != delegation:
                call.delegation = delegation

            # update zipcode
            if call.zipcode != zipcode:
                call.zipcode = zipcode

            # update city
            if call.city != city:
                call.city = city

            # update state
            if call.state != state:
                call.state = state

            # update location
            if call.location.x != longitude or call.location.y != latitude:
                call.location = Point(longitude, latitude, srid=4326)

            # update assignment
            if call.assignment != assignment:
                call.assignment = assignment

            # update description
            if call.description != description:
                call.description = description

            # update call_time
            if call.call_time != call_time:
                call.call_time = call_time

            # update departure_time
            if call.departure_time != departure_time:
                call.departure_time = departure_time

            # update transfer_time
            if call.transfer_time != transfer_time:
                call.transfer_time = transfer_time

            # update hospital_time
            if call.hospital_time != hospital_time:
                call.hospital_time = hospital_time

            # update base_time
            if call.base_time != base_time:
                call.base_time = base_time

            call.save()
            self.stdout.write(self.style.SUCCESS(">> Updated call {} in database").format(call.id))

        # If object does not exist on database, create it
        except ObjectDoesNotExist:
            serializer = CallSerializer(data=data)

            if serializer.is_valid():
               try:
                    # Save the call to the db with the serializer
                   call = serializer.save()
                   self.stdout.write(self.style.SUCCESS(">> Wrote new call {} to database").format(call.id))
               except Exception:
                   self.stdout.write(self.style.ERROR("*> Error writing Call to database"))
            else:
                self.stdout.write("Serializer Not Valid" + str(data))

        except Exception as e:
            self.stdout.write(self.style.ERROR("*> Error with parsing call fields"))

    # Update user location
    def on_user_loc(self, client, userdata, msg):

        topic = msg.topic.split('/')
        user = topic[1]

        if not msg.payload:
            return

        # Parse data into json dict
        stream = BytesIO(msg.payload)
        data = JSONParser().parse(stream)

        # TODO Find out which ambulance is linked to user
        ambulance = 1
        data['ambulance'] = ambulance

        # Serialize data into object
        serializer = MQTTLocationSerializer(data=data)

        if serializer.is_valid():
            try:
                # MAURICIO: Does that save the location to the database? Are we saving the user location as well?
                lp = serializer.save()

                # Publish ambulance location as soon as user location saved
                try:
                    self.pub_ambulance_loc(client, ambulance, lp)
                except Exception as e:
                    self.stdout.write(
                        self.style.ERROR("*> Failed to publish location. Exception: {}".format(e)))

                self.stdout.write(
                        self.style.SUCCESS(">> LocationPoint for user {} in ambulance {} successfully created.".format(user, ambulance)))
            except Exception as e:
                self.stdout.write(
                    self.style.ERROR("*> LocationPoint for user {} in ambulance {} failed to create. Exception: {}".format(user, ambulance, e)))

        else:
            self.stdout.write(
                self.style.ERROR("*> Input data for user location invalid"))

    # Publish location for ambulance - after ambulance starts querying location table,
    # take out the dependency on the lp generated
    def pub_ambulance_loc(self, client, amb_id, lp):
        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(">> Publishing location for ambulance {}".format(amb_id)))

        try:
            ambulance = Ambulances.objects.get(id=amb_id)

            # Set status and grab latest location; status should be calculated
            # rather than what this code is doing
            # status = Status.objects.all().first()
            # ambulance.status = status

            # Calculate ambulance orientation?
            ambulance.orientation = 200

            # Save changes made to ambulance
            ambulance.save()

            # Convert obj back to json
            serializer = MQTTAmbulanceLocSerializer(ambulance)
            json = JSONRenderer().render(serializer.data)

            # Publish json - be sure to do this in the seeder
            client.publish('ambulance/{}/location'.format(amb_id), json, qos=2, retain=True)
            client.publish('ambulance/{}/status'.format(amb_id), ambulance.status.name, qos=2, retain=True)

        except ObjectDoesNotExist:
            self.stdout.write(
                self.style.ERROR("*> Ambulance {} does not exist".format(amb_id)))

        except Exception as e:
            print(e)
            self.stdout.write(
                self.style.ERROR("Error parsing data"))

    # Update status from dispatch team
    def on_status(self, client, userdata, msg):
        topic = msg.topic.split('/')
        amb_id = topic[1]

        if not msg.payload:
            return

        status_str = msg.payload.decode("utf-8")
        ambulance = None

        try:
            ambulance = Ambulances.objects.get(id=amb_id)

        except ObjectDoesNotExist:
            self.stdout.write(
                self.style.ERROR("*> Ambulance {} does not exist".format(amb_id)))
            return

        try:
            status = Status.objects.get(name=status_str)
            if ambulance.status.name != status_str:
                ambulance.status = status
                ambulance.save()
            self.stdout.write(self.style.SUCCESS(
                ">> Successful status update: {} for ambulance {}").format(status_str, amb_id))

        except ObjectDoesNotExist:
            self.stdout.write(
                self.style.ERROR("*> Status {} does not exist".format(status_str)))
            return

        except Exception as e:
            print(e)
            self.stdout.write(
                self.style.ERROR("*> Error saving status for ambulance {}".format(amb_id)))

class Command(BaseCommand):
    help = 'Connect to the mqtt broker'

    def handle(self, *args, **options):

        broker = {
            'USERNAME': '',
            'PASSWORD': '',
            'HOST': 'localhost',
            'PORT': 1883,
            'KEEPALIVE': 60,
            'CLIENT_ID': 'django',
            'CLEAN_SESSION': True
        }
        broker.update(settings.MQTT)

        client = Client(broker, self.stdout, self.style,
                        verbosity = options['verbosity'])

        self.stdout.write(
            self.style.SUCCESS("""* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
* * *                    M Q T T   C L I E N T                    * * *
* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"""))

        try:
            client.loop_forever()

        except KeyboardInterrupt:
            pass

        finally:
            client.disconnect()
<|endoftext|>


<commit_before>def arg(x):
    """
    The argument function.
    """
    from cmath import log
    return log(x).imag

def generate_values():
    """
    Create values to test the function at.
    """
    from math import sin, cos, pi
    from random import random
    # Generate 3 circles in complex plane, with diameters 0.5, 1 and 2. We
    # avoid special values like -1, +/- i, etc., because they typically send
    # the numerical values close to the branch cut, and numerical errors then
    # flip the sign, e.g.:
    #    In [1]: sqrt((-0.5j)**2)
    #    Out[1]: -0.5j
    #
    #    In [2]: (-0.5j)**2
    #    Out[2]: (-0.25-0j)
    #
    #    In [3]: sqrt(-0.25)
    #    Out[3]: 0.5j
    # Here [3] is the correct value and [1] is incorrect, but that happens due
    # to the round off errors in [2] (the small negative imaginary part makes
    # sqrt() return -0.5j instead of +0.5j).
    #
    # For this reason, we chose N=7.
    N = 7
    circle = []
    for n in range(N): circle.append(cos(2*pi*n/N)+1j*sin(2*pi*n/N))
    values = []
    for n in range(N): values.append(0.5*circle[n])
    for n in range(N): values.append(1.0*circle[n])
    for n in range(N): values.append(2.0*circle[n])
    # Add some random points:
    for n in range(30):
        values.append((random()-0.5)*20 + 1j*(random()-0.5)*20)
    return values
values = generate_values()

def feq(a, b, max_relative_error=1e-12, max_absolute_error=1e-12):
    """
    Returns True if a==b to the given relative and absolute errors, otherwise
    False.
    """
    # if the numbers are close enough (absolutely), then they are equal
    if abs(a-b) < max_absolute_error:
        return True
    # if not, they can still be equal if their relative error is small
    if abs(b) > abs(a):
        relative_error = abs(a-b)/abs(b)
    else:
        relative_error = abs(a-b)/abs(a)
    #print abs(a-b), relative_error
    return relative_error <= max_relative_error

def test_zero1(lhs, rhs):
    """
    Tests that a complex function f(x) of one complex variable is zero.
    """
    for x in values:
        r = feq(lhs(x), rhs(x))
        if not r:
            print "x lhs(x) rhs(x)"
            print x, lhs(x), rhs(x)
            assert False

def test_zero2(lhs, rhs):
    """
    Tests that a complex function f(x, y) of two complex variables is zero.
    """
    for x in values:
        for y in values:
            r = feq(lhs(x, y), rhs(x, y))
            if not r:
                print "x y lhs(x, y) rhs(x, y)"
                print x, y, lhs(x, y), rhs(x, y)
                assert False

def test_zero3(lhs, rhs):
    """
    Tests that a complex function f(x, y, z) of three complex variables is zero.
    """
    for x in values:
        for y in values:
            for z in values:
                r = feq(lhs(x, y, z), rhs(x, y, z))
                if not r:
                    print "x y z lhs(x, y, z) rhs(x, y, z)"
                    print x, y, z, lhs(x, y, z), rhs(x, y, z)
                    assert False

def diff(f, z0, theta, eps=1e-8):
    h = eps*exp(I*theta)
    return (f(z0+h)-f(z0)) / h

def diff2(dfdz, dfdconjz, z0, theta):
    return dfdz(z0) + dfdconjz(z0)*exp(-2*I*theta)

def test_zero(f, dfdz, dfdconjz, z0, theta, eps=1e-8):
    assert feq(diff(f, z0, theta, eps), diff2(dfdz, dfdconjz, z0, theta),
            max_relative_error=eps*1e2, max_absolute_error=eps*1e2)

from math import floor, pi
from cmath import sqrt, exp, log
I = 1j

angles = [0, pi/7, pi/4, pi/2, 3*pi/4, pi]

for x in values:
    for theta in angles:
        test_zero(lambda x: abs(x), lambda x: x.conjugate()/(2*abs(x)),
                lambda x: x/(2*abs(x)), x, theta)
        test_zero(lambda x: log(x), lambda x: 1/x, lambda x: 0, x, theta)
        test_zero(lambda x: log(exp(x-x.conjugate())), lambda x: 1,
                lambda x: -1, x, theta)
<commit_msg>Add two more tests<commit_after>def arg(x):
    """
    The argument function.
    """
    from cmath import log
    return log(x).imag

def generate_values():
    """
    Create values to test the function at.
    """
    from math import sin, cos, pi
    from random import random
    # Generate 3 circles in complex plane, with diameters 0.5, 1 and 2. We
    # avoid special values like -1, +/- i, etc., because they typically send
    # the numerical values close to the branch cut, and numerical errors then
    # flip the sign, e.g.:
    #    In [1]: sqrt((-0.5j)**2)
    #    Out[1]: -0.5j
    #
    #    In [2]: (-0.5j)**2
    #    Out[2]: (-0.25-0j)
    #
    #    In [3]: sqrt(-0.25)
    #    Out[3]: 0.5j
    # Here [3] is the correct value and [1] is incorrect, but that happens due
    # to the round off errors in [2] (the small negative imaginary part makes
    # sqrt() return -0.5j instead of +0.5j).
    #
    # For this reason, we chose N=7.
    N = 7
    circle = []
    for n in range(N): circle.append(cos(2*pi*n/N)+1j*sin(2*pi*n/N))
    values = []
    for n in range(N): values.append(0.5*circle[n])
    for n in range(N): values.append(1.0*circle[n])
    for n in range(N): values.append(2.0*circle[n])
    # Add some random points:
    for n in range(30):
        values.append((random()-0.5)*20 + 1j*(random()-0.5)*20)
    return values
values = generate_values()

def feq(a, b, max_relative_error=1e-12, max_absolute_error=1e-12):
    """
    Returns True if a==b to the given relative and absolute errors, otherwise
    False.
    """
    # if the numbers are close enough (absolutely), then they are equal
    if abs(a-b) < max_absolute_error:
        return True
    # if not, they can still be equal if their relative error is small
    if abs(b) > abs(a):
        relative_error = abs(a-b)/abs(b)
    else:
        relative_error = abs(a-b)/abs(a)
    #print abs(a-b), relative_error
    return relative_error <= max_relative_error

def test_zero1(lhs, rhs):
    """
    Tests that a complex function f(x) of one complex variable is zero.
    """
    for x in values:
        r = feq(lhs(x), rhs(x))
        if not r:
            print "x lhs(x) rhs(x)"
            print x, lhs(x), rhs(x)
            assert False

def test_zero2(lhs, rhs):
    """
    Tests that a complex function f(x, y) of two complex variables is zero.
    """
    for x in values:
        for y in values:
            r = feq(lhs(x, y), rhs(x, y))
            if not r:
                print "x y lhs(x, y) rhs(x, y)"
                print x, y, lhs(x, y), rhs(x, y)
                assert False

def test_zero3(lhs, rhs):
    """
    Tests that a complex function f(x, y, z) of three complex variables is zero.
    """
    for x in values:
        for y in values:
            for z in values:
                r = feq(lhs(x, y, z), rhs(x, y, z))
                if not r:
                    print "x y z lhs(x, y, z) rhs(x, y, z)"
                    print x, y, z, lhs(x, y, z), rhs(x, y, z)
                    assert False

def diff(f, z0, theta, eps=1e-8):
    h = eps*exp(I*theta)
    return (f(z0+h)-f(z0)) / h

def diff2(dfdz, dfdconjz, z0, theta):
    return dfdz(z0) + dfdconjz(z0)*exp(-2*I*theta)

def test_zero(f, dfdz, dfdconjz, z0, theta, eps=1e-8):
    assert feq(diff(f, z0, theta, eps), diff2(dfdz, dfdconjz, z0, theta),
            max_relative_error=eps*1e2, max_absolute_error=eps*1e2)

from math import floor, pi
from cmath import sqrt, exp, log
I = 1j

angles = [0, pi/7, pi/4, pi/2, 3*pi/4, pi]

for x in values:
    for theta in angles:
        test_zero(lambda x: abs(x), lambda x: x.conjugate()/(2*abs(x)),
                lambda x: x/(2*abs(x)), x, theta)
        test_zero(lambda x: log(x), lambda x: 1/x, lambda x: 0, x, theta)
        test_zero(lambda x: log(exp(x-x.conjugate())), lambda x: 1,
                lambda x: -1, x, theta)
        test_zero(lambda x: sqrt(x**2), lambda x: sqrt(x**2)/x, lambda x: 0, x, theta)
        test_zero(lambda x: sqrt(x**2), lambda x: x/sqrt(x**2), lambda x: 0, x, theta)
<|endoftext|>


<commit_before>import re
import logging
from anchore_engine.services.policy_engine.engine.policy.gate import BaseTrigger, Gate
from anchore_engine.services.policy_engine.engine.policy.utils import RegexParamValidator, CommaDelimitedNumberListValidator, NameVersionListValidator, \
    CommaDelimitedStringListValidator, delim_parser, TypeValidator, InputValidator
from anchore_engine.services.policy_engine.engine.logs import get_logger

log = get_logger()

class DockerfileDirectiveListValidator(InputValidator):
    __directives__ = [
        'ADD',
        'ARG',
        'LABEL',
        'COPY',
        'CMD',
        'ENTRYPOINT',
        'ENV',
        'EXPOSE',
        'FROM',
        'HEALTHCHECK',
        'LABEL',
        'MAINTAINER',
        'ONBUILD',
        'USER',
        'RUN',
        'SHELL',
        'STOPSIGNAL',
        'VOLUME',
        'WORKDIR'
    ]

    def validation_criteria(self):
        return 'In: {}'.format(self.__directives__)

    def __call__(self, *args, **kwargs):
        if args and args[0]:
            parts = map(lambda x: x.strip(), args[0].split(','))
            return not bool(filter(lambda x: x not in self.__directives__, parts))
        else:
            return False

class ConditionValidator(InputValidator):
    __conditions__ = [
        '=',
        '!=',
        'exists',
        'not_exists',
        'like',
        'not_like'
    ]

    def validation_criteria(self):
        return 'In: {}'.format(self.__conditions__)

    def __call__(self, *args, **kwargs):
        if args and args[0]:
            return args[0].strip() in self.__conditions__
        return False

class EffectiveUserTrigger(BaseTrigger):
    __trigger_name__ = 'EFFECTIVEUSER'
    __description__ = 'Triggers if the effective user for the container is either root when not allowed or is not in a whitelist'

    __params__ = {
        'ALLOWED': CommaDelimitedStringListValidator(),
        'DENIED': CommaDelimitedStringListValidator(),
    }

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return # Prep step blocked this eval due to condition on the dockerfile, so skip

        allowed_users = delim_parser(self.eval_params.get('ALLOWED', ''))
        denied_users = delim_parser(self.eval_params.get('DENIED', ''))
        user_lines = context.data.get('prepared_dockerfile').get('USER', [])

        # If not overt, make it so
        if not user_lines:
            user_lines = ['USER root']

        user = user_lines[-1].strip()  # The last USER line is the determining entry
        user = user[len('USER'):].strip()
        if allowed_users and user not in allowed_users:
            self._fire(msg='User {} found as effective user, which is not on the allowed list'.format(user))
        if denied_users and user in denied_users:
            self._fire(msg='User {} found as effective user, which is on the denied list'.format(user))

class DirectiveCheckTrigger(BaseTrigger):
    __trigger_name__ = 'DIRECTIVECHECK'
    __description__ = 'Triggers if any directives in the list are found to match the described condition in the dockerfile'

    __params__ = {
        'DIRECTIVES': DockerfileDirectiveListValidator(),
        'CHECK': ConditionValidator(),
        'CHECK_VALUE': TypeValidator(str)
    }

    _conditions_requiring_check_val = [
        '=', '!=', 'like', 'not_like'
    ]

    ops = {
        '=': lambda x, y: x == y,
        '!=': lambda x, y: x != y,
        'exists': lambda x, y: True,
        'not_exists': lambda x, y: False,
        'like': lambda x, y: bool(re.match(y, x)),
        'not_like': lambda x, y: not bool(re.match(y, x))
    }

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return # Prep step blocked this eval due to condition on the dockerfile, so skip

        directives = set(map(lambda x: x.upper(), delim_parser(self.eval_params.get('DIRECTIVES', ''))))
        condition = self.eval_params.get('CHECK')
        check_value = self.eval_params.get('CHECK_VALUE')
        operation = self.ops[condition]

        if not condition or not directives:
            return

        df = context.data.get('prepared_dockerfile')

        for directive, lines in filter(lambda x: x[0] in directives, df.items()):
            for l in lines:
                l = l[len(directive):].strip()
                if operation(l, check_value):
                    self._fire(msg="Dockerfile directive '{}' check '{}' matched against '{}' for line '{}'".format(directive, condition, check_value if check_value else '', l))

        if condition == 'not_exists':
            for match in directives.difference(directives.intersection(set(map(lambda x: x.upper(), df.keys())))):
                self._fire(msg="Dockerfile directive '{}' not found, matching condition '{}' check".format(match, condition))

class ExposeTrigger(BaseTrigger):
    __trigger_name__ = 'EXPOSE'

    __params__ = {
        'ALLOWEDPORTS': CommaDelimitedNumberListValidator(),
        'DENIEDPORTS': CommaDelimitedNumberListValidator()
    }
    __description__ = 'triggers if Dockerfile is EXPOSEing ports that are not in ALLOWEDPORTS, or are in DENIEDPORTS'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        allowed_ports = delim_parser(self.eval_params.get('ALLOWEDPORTS', ''))
        denied_ports = delim_parser(self.eval_params.get('DENIEDPORTS', ''))

        expose_lines = context.data.get('prepared_dockerfile', {}).get('EXPOSE', [])
        for line in expose_lines:
            matchstr = None
            line = line.strip()
            if re.match("^\s*(EXPOSE|" + 'EXPOSE'.lower() + ")\s+(.*)", line):
                matchstr = re.match("^\s*(EXPOSE|" + 'EXPOSE'.lower() + ")\s+(.*)", line).group(2)

            if matchstr:
                iexpose = matchstr.split()
                if denied_ports:
                    if 'ALL' in denied_ports and len(iexpose) > 0:
                        self._fire(msg="Dockerfile exposes network ports but policy sets DENIEDPORTS=ALL: " + str(iexpose))
                    else:
                        for p in denied_ports:
                            if p in iexpose:
                                self._fire(msg="Dockerfile exposes port (" + p + ") which is in policy file DENIEDPORTS list")
                            elif p + '/tcp' in iexpose:
                                self._fire(msg="Dockerfile exposes port (" + p + "/tcp) which is in policy file DENIEDPORTS list")
                            elif p + '/udp' in iexpose:
                                self._fire(msg="Dockerfile exposes port (" + p + "/udp) which is in policy file DENIEDPORTS list")

                if allowed_ports:
                    if 'NONE' in allowed_ports and len(iexpose) > 0:
                        self._fire(msg="Dockerfile exposes network ports but policy sets ALLOWEDPORTS=NONE: " + str(iexpose))
                    else:
                        for p in allowed_ports:
                            done = False
                            while not done:
                                try:
                                    iexpose.remove(p)
                                    done = False
                                except:
                                    done = True

                                try:
                                    iexpose.remove(p + '/tcp')
                                    done = False
                                except:
                                    done = True

                                try:
                                    iexpose.remove(p + '/udp')
                                    done = False
                                except:
                                    done = True

                        for ip in iexpose:
                            self._fire(msg="Dockerfile exposes port (" + ip + ") which is not in policy file ALLOWEDPORTS list")

                        # Replaecable by:
                        # for port in filter(lambda x: x.split('/')[0] not in allowed_ports, iexpose):
                        #   self._fire(...)
        return

class NoFromTrigger(BaseTrigger):
    __trigger_name__ = 'NOFROM'
    __params__ = None
    __description__ = 'triggers if there is no FROM line specified in the Dockerfile'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        from_lines = context.data['prepared_dockerfile'].get('FROM')
        if not from_lines:
            self._fire(msg="No 'FROM' directive in Dockerfile")
            return

class FromScratch(BaseTrigger):
    __trigger_name__ = 'FROMSCRATCH'
    __description__ = 'triggers the FROM line specified "scratch" as the parent'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        from_lines = context.data['prepared_dockerfile'].get('FROM', [])
        for line in from_lines:
            fromstr = None
            if re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line):
                fromstr = re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line).group(2)

            if fromstr == 'SCRATCH' or fromstr.lower() == 'scratch':
                self._fire(msg="'FROM' container is 'scratch' - (" + str(fromstr) + ")")

class NoTag(BaseTrigger):
    __trigger_name__ = 'NOTAG'
    __description__ = 'triggers if the FROM container specifies a repo but no explicit, non-latest tag'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        from_lines = context.data['prepared_dockerfile'].get('FROM', [])
        for line in from_lines:
            fromstr = None
            if re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line):
                fromstr = re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line).group(2)

            if fromstr:
                if re.match("(\S+):(\S+)", fromstr):
                    repo, tag = re.match("(\S+):(\S+)", fromstr).group(1, 2)
                    if tag == 'latest':
                        self._fire(msg="container does not specify a non-latest container tag - (" + str(
                            fromstr) + ")")
                else:
                    self._fire(msg="container does not specify a non-latest container tag - (" + str(fromstr) + ")")

class Sudo(BaseTrigger):
    __trigger_name__ = 'SUDO'
    __description__ = 'triggers if the Dockerfile contains operations running with sudo'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        if image_obj.dockerfile_contents:
            for line in image_obj.dockerfile_contents.splitlines():
                line = line.strip()
                if re.match(".*sudo.*", line):
                    self._fire(msg="Dockerfile contains a 'sudo' command: " + str(line))

class VolumePresent(BaseTrigger):
    __trigger_name__ = 'VOLUMEPRESENT'
    __description__ = 'triggers if the Dockerfile contains a VOLUME line'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        for line in context.data['prepared_dockerfile'].get('VOLUME', []):
            self._fire(msg='Dockerfile contains a VOLUME line: ' + str(line))

class NoHealthCheck(BaseTrigger):
    __trigger_name__ = 'NOHEALTHCHECK'
    __description__ = 'triggers if the Dockerfile does not contain any HEALTHCHECK instructions'
    __msg__ = 'Dockerfile does not contain any HEALTHCHECK instructions'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        if not context.data['prepared_dockerfile'].get('HEALTHCHECK'):
            self._fire()

class NoDockerfile(BaseTrigger):
    __trigger_name__ = 'NODOCKERFILE'
    __description__ = 'triggers if anchore analysis was performed without supplying a real Dockerfile'
    __msg__ = 'Image was not analyzed with an actual Dockerfile'

    def evaluate(self, image_obj, context):
        """
        Evaluate using the initialized values for this object:
        """
        if image_obj.dockerfile_mode != 'Actual':
            self._fire()

class DockerfileGate(Gate):
    __gate_name__ = 'DOCKERFILECHECK'

    __triggers__ = [
        DirectiveCheckTrigger,
        EffectiveUserTrigger,
        ExposeTrigger,
        NoFromTrigger,
        FromScratch,
        NoTag,
        Sudo,
        VolumePresent,
        NoHealthCheck,
        NoDockerfile
    ]

    def prepare_context(self, image_obj, context):
        """
        Pre-processes the image's dockerfile.
        Leaves the context with a dictionary of dockerfile lines by directive.
        e.g.
        context.data['dockerfile']['RUN'] = ['RUN apt-get update', 'RUN blah']
        context.data['dockerfile']['VOLUME'] = ['VOLUME /tmp', 'VOLUMN /var/log']

        :return: updated context
        """

        # Optimization by single-pass parsing the docker file instead of one per trigger eval.
        # unknown/known is up to each trigger

        if image_obj.dockerfile_mode == "Unknown":
            return

        context.data['prepared_dockerfile'] = {}

        if image_obj.dockerfile_contents:
            linebuf = ""
            for line in image_obj.dockerfile_contents.splitlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    patt = re.match(".*\\\$", line)
                    if patt:
                        line = re.sub("\\\$", "", line)
                        linebuf = linebuf + line
                    else:
                        linebuf = linebuf + line
                        if linebuf:
                            directive,remainder = linebuf.split(' ', 1)
                            directive = directive.upper()
                            if directive not in context.data['prepared_dockerfile']:
                                context.data['prepared_dockerfile'][directive] = []
                            context.data['prepared_dockerfile'][directive].append(linebuf)
                            linebuf = ""
                else:
                    continue
                    # Skip comment lines in the dockerfile

            #for line in image_obj.dockerfile_contents.splitlines():
            #    line = line.strip()
            #    if line and not line.startswith('#'):
            #        directive, remainder = line.split(' ', 1)
            #        directive = directive.upper() # Normalize to all upper
            #        if directive not in context.data['prepared_dockerfile']:
            #            context.data['prepared_dockerfile'][directive] = []
            #        context.data['prepared_dockerfile'][directive].append(line)
            #    else:
            #        continue
            #        # Skip comment lines in the dockerfile

        return context
<commit_msg>cleanup fixes #3<commit_after>import re
import logging
from anchore_engine.services.policy_engine.engine.policy.gate import BaseTrigger, Gate
from anchore_engine.services.policy_engine.engine.policy.utils import RegexParamValidator, CommaDelimitedNumberListValidator, NameVersionListValidator, \
    CommaDelimitedStringListValidator, delim_parser, TypeValidator, InputValidator
from anchore_engine.services.policy_engine.engine.logs import get_logger

log = get_logger()

class DockerfileDirectiveListValidator(InputValidator):
    __directives__ = [
        'ADD',
        'ARG',
        'LABEL',
        'COPY',
        'CMD',
        'ENTRYPOINT',
        'ENV',
        'EXPOSE',
        'FROM',
        'HEALTHCHECK',
        'LABEL',
        'MAINTAINER',
        'ONBUILD',
        'USER',
        'RUN',
        'SHELL',
        'STOPSIGNAL',
        'VOLUME',
        'WORKDIR'
    ]

    def validation_criteria(self):
        return 'In: {}'.format(self.__directives__)

    def __call__(self, *args, **kwargs):
        if args and args[0]:
            parts = map(lambda x: x.strip(), args[0].split(','))
            return not bool(filter(lambda x: x not in self.__directives__, parts))
        else:
            return False

class ConditionValidator(InputValidator):
    __conditions__ = [
        '=',
        '!=',
        'exists',
        'not_exists',
        'like',
        'not_like'
    ]

    def validation_criteria(self):
        return 'In: {}'.format(self.__conditions__)

    def __call__(self, *args, **kwargs):
        if args and args[0]:
            return args[0].strip() in self.__conditions__
        return False

class EffectiveUserTrigger(BaseTrigger):
    __trigger_name__ = 'EFFECTIVEUSER'
    __description__ = 'Triggers if the effective user for the container is either root when not allowed or is not in a whitelist'

    __params__ = {
        'ALLOWED': CommaDelimitedStringListValidator(),
        'DENIED': CommaDelimitedStringListValidator(),
    }

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return # Prep step blocked this eval due to condition on the dockerfile, so skip

        allowed_users = delim_parser(self.eval_params.get('ALLOWED', ''))
        denied_users = delim_parser(self.eval_params.get('DENIED', ''))
        user_lines = context.data.get('prepared_dockerfile').get('USER', [])

        # If not overt, make it so
        if not user_lines:
            user_lines = ['USER root']

        user = user_lines[-1].strip()  # The last USER line is the determining entry
        user = user[len('USER'):].strip()
        if allowed_users and user not in allowed_users:
            self._fire(msg='User {} found as effective user, which is not on the allowed list'.format(user))
        if denied_users and user in denied_users:
            self._fire(msg='User {} found as effective user, which is on the denied list'.format(user))

class DirectiveCheckTrigger(BaseTrigger):
    __trigger_name__ = 'DIRECTIVECHECK'
    __description__ = 'Triggers if any directives in the list are found to match the described condition in the dockerfile'

    __params__ = {
        'DIRECTIVES': DockerfileDirectiveListValidator(),
        'CHECK': ConditionValidator(),
        'CHECK_VALUE': TypeValidator(str)
    }

    _conditions_requiring_check_val = [
        '=', '!=', 'like', 'not_like'
    ]

    ops = {
        '=': lambda x, y: x == y,
        '!=': lambda x, y: x != y,
        'exists': lambda x, y: True,
        'not_exists': lambda x, y: False,
        'like': lambda x, y: bool(re.match(y, x)),
        'not_like': lambda x, y: not bool(re.match(y, x))
    }

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return # Prep step blocked this eval due to condition on the dockerfile, so skip

        directives = set(map(lambda x: x.upper(), delim_parser(self.eval_params.get('DIRECTIVES', ''))))
        condition = self.eval_params.get('CHECK')
        check_value = self.eval_params.get('CHECK_VALUE')
        operation = self.ops[condition]

        if not condition or not directives:
            return

        df = context.data.get('prepared_dockerfile')

        for directive, lines in filter(lambda x: x[0] in directives, df.items()):
            for l in lines:
                l = l[len(directive):].strip()
                if operation(l, check_value):
                    self._fire(msg="Dockerfile directive '{}' check '{}' matched against '{}' for line '{}'".format(directive, condition, check_value if check_value else '', l))

        if condition == 'not_exists':
            for match in directives.difference(directives.intersection(set(map(lambda x: x.upper(), df.keys())))):
                self._fire(msg="Dockerfile directive '{}' not found, matching condition '{}' check".format(match, condition))

class ExposeTrigger(BaseTrigger):
    __trigger_name__ = 'EXPOSE'

    __params__ = {
        'ALLOWEDPORTS': CommaDelimitedNumberListValidator(),
        'DENIEDPORTS': CommaDelimitedNumberListValidator()
    }
    __description__ = 'triggers if Dockerfile is EXPOSEing ports that are not in ALLOWEDPORTS, or are in DENIEDPORTS'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        allowed_ports = delim_parser(self.eval_params.get('ALLOWEDPORTS', ''))
        denied_ports = delim_parser(self.eval_params.get('DENIEDPORTS', ''))

        expose_lines = context.data.get('prepared_dockerfile', {}).get('EXPOSE', [])
        for line in expose_lines:
            matchstr = None
            line = line.strip()
            if re.match("^\s*(EXPOSE|" + 'EXPOSE'.lower() + ")\s+(.*)", line):
                matchstr = re.match("^\s*(EXPOSE|" + 'EXPOSE'.lower() + ")\s+(.*)", line).group(2)

            if matchstr:
                iexpose = matchstr.split()
                if denied_ports:
                    if 'ALL' in denied_ports and len(iexpose) > 0:
                        self._fire(msg="Dockerfile exposes network ports but policy sets DENIEDPORTS=ALL: " + str(iexpose))
                    else:
                        for p in denied_ports:
                            if p in iexpose:
                                self._fire(msg="Dockerfile exposes port (" + p + ") which is in policy file DENIEDPORTS list")
                            elif p + '/tcp' in iexpose:
                                self._fire(msg="Dockerfile exposes port (" + p + "/tcp) which is in policy file DENIEDPORTS list")
                            elif p + '/udp' in iexpose:
                                self._fire(msg="Dockerfile exposes port (" + p + "/udp) which is in policy file DENIEDPORTS list")

                if allowed_ports:
                    if 'NONE' in allowed_ports and len(iexpose) > 0:
                        self._fire(msg="Dockerfile exposes network ports but policy sets ALLOWEDPORTS=NONE: " + str(iexpose))
                    else:
                        for p in allowed_ports:
                            done = False
                            while not done:
                                try:
                                    iexpose.remove(p)
                                    done = False
                                except:
                                    done = True

                                try:
                                    iexpose.remove(p + '/tcp')
                                    done = False
                                except:
                                    done = True

                                try:
                                    iexpose.remove(p + '/udp')
                                    done = False
                                except:
                                    done = True

                        for ip in iexpose:
                            self._fire(msg="Dockerfile exposes port (" + ip + ") which is not in policy file ALLOWEDPORTS list")

                        # Replaecable by:
                        # for port in filter(lambda x: x.split('/')[0] not in allowed_ports, iexpose):
                        #   self._fire(...)
        return

class NoFromTrigger(BaseTrigger):
    __trigger_name__ = 'NOFROM'
    __params__ = None
    __description__ = 'triggers if there is no FROM line specified in the Dockerfile'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        from_lines = context.data['prepared_dockerfile'].get('FROM')
        if not from_lines:
            self._fire(msg="No 'FROM' directive in Dockerfile")
            return

class FromScratch(BaseTrigger):
    __trigger_name__ = 'FROMSCRATCH'
    __description__ = 'triggers the FROM line specified "scratch" as the parent'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        from_lines = context.data['prepared_dockerfile'].get('FROM', [])
        for line in from_lines:
            fromstr = None
            if re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line):
                fromstr = re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line).group(2)

            if fromstr == 'SCRATCH' or fromstr.lower() == 'scratch':
                self._fire(msg="'FROM' container is 'scratch' - (" + str(fromstr) + ")")

class NoTag(BaseTrigger):
    __trigger_name__ = 'NOTAG'
    __description__ = 'triggers if the FROM container specifies a repo but no explicit, non-latest tag'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        from_lines = context.data['prepared_dockerfile'].get('FROM', [])
        for line in from_lines:
            fromstr = None
            if re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line):
                fromstr = re.match("^\s*(FROM|" + 'FROM'.lower() + ")\s+(.*)", line).group(2)

            if fromstr:
                if re.match("(\S+):(\S+)", fromstr):
                    repo, tag = re.match("(\S+):(\S+)", fromstr).group(1, 2)
                    if tag == 'latest':
                        self._fire(msg="container does not specify a non-latest container tag - (" + str(
                            fromstr) + ")")
                else:
                    self._fire(msg="container does not specify a non-latest container tag - (" + str(fromstr) + ")")

class Sudo(BaseTrigger):
    __trigger_name__ = 'SUDO'
    __description__ = 'triggers if the Dockerfile contains operations running with sudo'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        if image_obj.dockerfile_contents:
            for line in image_obj.dockerfile_contents.splitlines():
                line = line.strip()
                if re.match(".*sudo.*", line):
                    self._fire(msg="Dockerfile contains a 'sudo' command: " + str(line))

class VolumePresent(BaseTrigger):
    __trigger_name__ = 'VOLUMEPRESENT'
    __description__ = 'triggers if the Dockerfile contains a VOLUME line'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        for line in context.data['prepared_dockerfile'].get('VOLUME', []):
            self._fire(msg='Dockerfile contains a VOLUME line: ' + str(line))

class NoHealthCheck(BaseTrigger):
    __trigger_name__ = 'NOHEALTHCHECK'
    __description__ = 'triggers if the Dockerfile does not contain any HEALTHCHECK instructions'
    __msg__ = 'Dockerfile does not contain any HEALTHCHECK instructions'

    def evaluate(self, image_obj, context):
        if not context.data.get('prepared_dockerfile'):
            return  # Prep step blocked this due to condition on the dockerfile, so skip

        if not context.data['prepared_dockerfile'].get('HEALTHCHECK'):
            self._fire()

class NoDockerfile(BaseTrigger):
    __trigger_name__ = 'NODOCKERFILE'
    __description__ = 'triggers if anchore analysis was performed without supplying a real Dockerfile'
    __msg__ = 'Image was not analyzed with an actual Dockerfile'

    def evaluate(self, image_obj, context):
        """
        Evaluate using the initialized values for this object:
        """
        if image_obj.dockerfile_mode != 'Actual':
            self._fire()

class DockerfileGate(Gate):
    __gate_name__ = 'DOCKERFILECHECK'

    __triggers__ = [
        DirectiveCheckTrigger,
        EffectiveUserTrigger,
        ExposeTrigger,
        NoFromTrigger,
        FromScratch,
        NoTag,
        Sudo,
        VolumePresent,
        NoHealthCheck,
        NoDockerfile
    ]

    def prepare_context(self, image_obj, context):
        """
        Pre-processes the image's dockerfile.
        Leaves the context with a dictionary of dockerfile lines by directive.
        e.g.
        context.data['dockerfile']['RUN'] = ['RUN apt-get update', 'RUN blah']
        context.data['dockerfile']['VOLUME'] = ['VOLUME /tmp', 'VOLUMN /var/log']

        :return: updated context
        """

        # Optimization by single-pass parsing the docker file instead of one per trigger eval.
        # unknown/known is up to each trigger

        if image_obj.dockerfile_mode == "Unknown":
            return

        context.data['prepared_dockerfile'] = {}

        if image_obj.dockerfile_contents:
            linebuf = ""
            for line in image_obj.dockerfile_contents.splitlines():
                line = line.strip()
                if line and not line.startswith('#'):
                    patt = re.match(".*\\\$", line)
                    if patt:
                        line = re.sub("\\\$", "", line)
                        linebuf = linebuf + line
                    else:
                        linebuf = linebuf + line
                        if linebuf:
                            directive,remainder = linebuf.split(' ', 1)
                            directive = directive.upper()
                            if directive not in context.data['prepared_dockerfile']:
                                context.data['prepared_dockerfile'][directive] = []
                            context.data['prepared_dockerfile'][directive].append(linebuf)
                            linebuf = ""
                else:
                    continue
                    # Skip comment lines in the dockerfile

        return context
<|endoftext|>


<commit_before>"""Peewee migrations -- 001_init_tables.py.

Some examples (model - class or model name)::

    > Model = migrator.orm['model_name']            # Return model in current state by name

    > migrator.sql(sql)                             # Run custom SQL
    > migrator.python(func, *args, **kwargs)        # Run python code
    > migrator.create_model(Model)                  # Create a model (could be used as decorator)
    > migrator.remove_model(model, cascade=True)    # Remove a model
    > migrator.add_fields(model, **fields)          # Add fields to a model
    > migrator.change_fields(model, **fields)       # Change fields
    > migrator.remove_fields(model, *field_names, cascade=True)
    > migrator.rename_field(model, old_field_name, new_field_name)
    > migrator.rename_table(model, new_table_name)
    > migrator.add_index(model, *col_names, unique=False)
    > migrator.drop_index(model, *col_names)
    > migrator.add_not_null(model, *field_names)
    > migrator.drop_not_null(model, *field_names)
    > migrator.add_default(model, field_name, default)

"""

import peewee as pw
import logging

from models import ThermostatGroup, Feature

try:
    import playhouse.postgres_ext as pw_pext
except ImportError:
    pass

SQL = pw.SQL

logger = logging.getLogger('openmotics')

def migrate(migrator, database, fake=False, **kwargs):
    ThermostatGroup.create(id=1, number=0, name='default', on=True)
    Feature.create(id=1, name='thermostats_gateway', enabled=False)

def rollback(migrator, database, fake=False, **kwargs):
    """Write your rollback migrations here."""
    pass
<commit_msg>make migrations idempotent<commit_after>"""Peewee migrations -- 001_init_tables.py.

Some examples (model - class or model name)::

    > Model = migrator.orm['model_name']            # Return model in current state by name

    > migrator.sql(sql)                             # Run custom SQL
    > migrator.python(func, *args, **kwargs)        # Run python code
    > migrator.create_model(Model)                  # Create a model (could be used as decorator)
    > migrator.remove_model(model, cascade=True)    # Remove a model
    > migrator.add_fields(model, **fields)          # Add fields to a model
    > migrator.change_fields(model, **fields)       # Change fields
    > migrator.remove_fields(model, *field_names, cascade=True)
    > migrator.rename_field(model, old_field_name, new_field_name)
    > migrator.rename_table(model, new_table_name)
    > migrator.add_index(model, *col_names, unique=False)
    > migrator.drop_index(model, *col_names)
    > migrator.add_not_null(model, *field_names)
    > migrator.drop_not_null(model, *field_names)
    > migrator.add_default(model, field_name, default)

"""

import peewee as pw
import logging

from models import ThermostatGroup, Feature

try:
    import playhouse.postgres_ext as pw_pext
except ImportError:
    pass

SQL = pw.SQL

logger = logging.getLogger('openmotics')

def migrate(migrator, database, fake=False, **kwargs):
    ThermostatGroup.get_or_create(number=0, name='default', on=True)
    Feature.get_or_create(name='thermostats_gateway', enabled=False)

def rollback(migrator, database, fake=False, **kwargs):
    """Write your rollback migrations here."""
    pass
<|endoftext|>


<commit_before>"""This is the HelloWorld challenge to teach the usage of the library by a
minimal example.

This challenge takes a number and a word and composes a new word based
on the input.

The class fully focuses on the "business logic" by overwriting just two
methods of parent class <challenge.challenge.Challenge>:

    * <build>: This sets up the data model from the given input lines.
    * <calc>: This is the core algorithm of the challenge.

This design pattern is called the "Template Method Pattern". While the parent
class handles the common parts of the challenge flow, the child class just
encapsulates what is special.

The two methods above is just the bare minimum to implement. For more complex
challenges you will overwrite or extend other parts of the parent class.

Have a look into the parent class, especially into the <main> function, that
controls the execution of the challenge. Also have a look into the parents
<__init__> method to see what instance attributes are already prepared to serve
the communication of the methods.

By the class attribute <sample> a minimal example of the possible input is
given. This is recommended for every challenge. It is useful during
development or for smoke tests. Otherwise this class attribute
will be overwritten by injection of an instance attribute of the same name.
"""

from challenges import Challenge

class HelloWorldChallenge(Challenge):
    """This is the HelloWorld Challenge class.

    This challenge takes a word as input and a number at which position
    to split the word. It returns a new word composed of the switched parts
    with a blank in between.
    """

    sample = '''
        5
        WorldHello
    '''

    """This is the input sample. It can be overwritten by injection.

        * first line: integer
        * second line: word

        It is used by the --klass option of the challenge runner.
        It can be used for unit tests.

        prompt> challenge HelloWorld/ --klass
        Hello World
    """

    expect = 'Hello World'

    """This is the expected output for the sample.

        It can be used for unit tests.
    """

    def build(self):
        """Parse the input lines and set up the model."""
        self.model.split_at = self.line_to_integers(0)[0]
        self.model.word = self.line(1)

    def calc(self):
        """Swap head and tail of the model and store to result."""
        first = self.model.word[self.model.split_at:]
        second = self.model.word[:self.model.split_at]
        self.result = first + ' ' + second
<commit_msg>Updat Hello World comments.<commit_after>"""This is the HelloWorld challenge to teach the usage of the library by a
minimal example.

This challenge takes a number and a word and composes a new word based
on the input.

The class fully focuses on the "business logic" by overwriting just two
methods of parent class <challenge.challenge.Challenge>:

    * <build>: This sets up the data model from the given input lines.
    * <calc>: This is the core algorithm of the challenge.

This design pattern is called the "Template Method Pattern". While the parent
class handles the common parts of the challenge flow, the child class just
encapsulates what is special.

The two methods above is just the bare minimum to implement. For more complex
challenges you will overwrite or extend other parts of the parent class.

Have a look into the parent class, especially into the <main> function, that
controls the execution of the challenge. Also have a look into the parents
<__init__> method to see what instance attributes are already prepared to serve
the communication of the methods.

By the class attribute <sample> a small example of the possible input is
given. This is recommended for every challenge. It is useful as input of
unit tests or for the smoke tests with the -k option. For more extended
input, in example from the given data files, this class attribute will be
overwritten by injection of an instance attribute of the same name.

The optional class attribute <expect> the expected output for the given
<sample> sets the expected output. This is typically used by unit tests.

Together <sample> and <expect> serve as documantation of the intention of
the challenge.
"""

from challenges import Challenge

class HelloWorldChallenge(Challenge):
    """This is the HelloWorld Challenge class.

    This challenge takes a word as input and a number at which position
    to split the word. It returns a new word composed of the switched parts
    with a blank in between.
    """

    sample = '''
        5
        WorldHello
    '''

    """This is a small exsample of the input. It can be overwritten by
    injection.

        It is used by the --klass option of the challenge runner.
        It is also used by unit tests.

        prompt> challenge HelloWorld/ --klass
        Hello World
    """

    expect = 'Hello World'

    """This is the expected output for the given sample.

        It is used by unit tests, especially by the full integration test.
    """

    def build(self):
        """Parse the input lines and set up the model."""
        self.model.split_at = self.line_to_integers(0)[0]
        self.model.word = self.line(1)

    def calc(self):
        """Swap head and tail of the model and store to result."""
        first = self.model.word[self.model.split_at:]
        second = self.model.word[:self.model.split_at]
        self.result = first + ' ' + second
<|endoftext|>