<commit_before>import numpy as np
import random
import sys

class insert_file_info :

    def __init__(self, full_file_path, filenumber, batch_size = 50,
        use_random_seed = False ) :
        """ full_file_path : full file path of the shuffled data
            filenumber     : An array of file number """
        self.filename         = full_file_path.rsplit('\\', 1)[-1]
        self.filename         = self.filename.rsplit('/', 1)[-1]
        self.full_file_path   = full_file_path
        self.nrows            = 0
        self.ncols            = 0
        self.nfile            = len(filenumber)
        self.batch_size       = batch_size
        self.current_index    = 0

    class DataSet(object) :
        file_info = None

        def __init__(self, images, labels, data_type = 'unknown') :
            #self.file_into = insert_file_info()

            #super(DataSet,self).__init__()
            #self.insert_file_info = insert_file_info

            self.data_type = data_type
            self.batch_size = 0
            self._ndata     = 0
            self._images = images
            self._labels = labels
            self._epochs_completed = 0
            self._index_in_epoch = 0
            self._index_in_datafile = 0
            self._file_index = 1

        #@staticmethod
        #def feed_self(self, batch_size, nrows) :
        #    self.batch_size = batch_size
        #    self.nrows      = nrows
            #print self.batch_size, self.nrows

        @property
        def images(self):
            return self._images

        @property
        def labels(self):
            return self._labels

        @property
        def ndata(self):
            return self._ndata

        @property
        def epochs_completed(self):
            return self._epochs_completed

        def next_batch(self, batch_size = 50, ndata = 8000) :
            self.batch_size = batch_size
            self._ndata     = ndata
            start = self._index_in_epoch
            if ( self._epochs_completed == 0 ) and ( start == 0 ) :
                print 'batch size : %d'    % self.batch_size
                print 'number of data: %d' % self._ndata
            self._index_in_epoch += self.batch_size
            if self._index_in_epoch > self._ndata :
                # Number of training epochs completed
                self._epochs_completed += 1
                # Shuffle data
                random.shuffle(self._images)
                random.shuffle(self._labels)
                # Reinitialize conunter
                start = 0
                self._index_in_epoch = self.batch_size
                assert self.batch_size <= self._ndata
            end = self._index_in_epoch
            return self._images[start:end], self._labels[start:end]

        def initialize_file_info_for_dose_of_data(self, full_file_path,
            filenumber, convert_test_labels_to_one_hot = True) :
            """ full_file_path : full file path of the shuffled data
                filenumber     : An array of file number """

            print 'Initializing file information for loading small dosage of data...'
            self.filename         = full_file_path.rsplit('\\', 1)[-1]
            self.filename         = self.filename.rsplit('/', 1)[-1]
            self.full_file_path   = full_file_path
            self.nfile            = len(filenumber)

            data = np.loadtxt(self.full_file_path%1)
            self.nrows, self.ncols = np.shape(data)
            self.nrows, self.ncols = int(self.nrows), int(self.ncols)

            # Use 10% of the data each for testing and validating, the remaining for
            # training
            self._nfile_train = int(self.nfile*.8)
            self._nfile_test  = int(self.nfile*.1)
            self._nfile_val   = self._nfile_test

            n_data_check = self.nfile - ( self._nfile_train + self._nfile_test + self._nfile_val )
            if n_data_check > 0 :
                self._nfile_train += n_data_check
            elif n_data_check < 0 :
                self._nfile_train -= n_data_check

            self._ndata           = self._nfile_train*self.nrows

            if self.data_type == 'train' :
                self.start_file_index = 1
                self.end_file_index = self._nfile_train
                self.convert_to_one_hot = True
            elif self.data_type == 'test' :
                self.start_file_index = self._nfile_train + 1
                self.end_file_index = self._nfile_train + self._nfile_test
                if convert_test_labels_to_one_hot :
                    self.convert_to_one_hot = convert_test_labels_to_one_hot
            elif self.data_type == 'validation' :
                self.start_file_index = self._nfile_train + self._nfile_test + 1
                self.end_file_index = self._nfile_train + self._nfile_test + self._nfile_validation
                self.convert_to_one_hot = False

        def next_dose(self, batch_size = 50) :

            def convert_to_one_hot( label ) :
                label_one_hot = np.zeros((len(label),2))
                for i in range(len(label)) :
                    label_one_hot[i,label[i]] = 1
                return label_one_hot

            start = self._index_in_datafile

            if ( self._file_index == self.start_file_index ) and ( start == 0 ) :
                self.batch_size = batch_size
                if np.modf(float(self.nrows)/self.batch_size)[0] > 0.0 :
                    self.batch_size = int(float(self.nrows)/20)

                print 'batch size : %d'    % self.batch_size
                print 'number of data: %d' % self._ndata

            self._index_in_datafile += self.batch_size
            if self._index_in_datafile > self.nrows :
                self._file_index += 1
                start = 0
                self._index_in_datafile = self.batch_size
                assert self.batch_size <= self.nrows

            if self._file_index > self.end_file_index :
                # Number of training epochs completed
                self._epochs_completed += 1
                self._file_index = self.start_file_index
                # Shufle data
                random.shuffle(self._images)
                random.shuffle(self._labels)
                # Reinitialize conunter
                start = 0
                self._index_in_datafile = self.batch_size

            end = self._index_in_datafile

            # Read in small dosage of data
            data = np.genfromtxt(self.full_file_path%(self._file_index) ,dtype=int,
                   skip_header=start, skip_footer=self.nrows-end)
            self._images = data[:,:-1].astype('int')
            labels = data[:,-1:].astype('int')
            if self.convert_to_one_hot :
                self._labels = convert_to_one_hot(labels)

            return self._images, self._labels

    def categorize_data(self, convert_test_labels_to_one_hot = True) :
        class DataSets(object):
            pass
        data_sets = DataSets()

        def convert_to_one_hot( label ) :
            label_one_hot = np.zeros((len(label),2))
            for i in range(len(label)) :
                label_one_hot[i,label[i]] = 1
            return label_one_hot

        data = np.loadtxt(self.full_file_path%1)
        self.nrows, self.ncols = np.shape(data)
        self.nrows, self.ncols = int(self.nrows), int(self.ncols)

        if np.modf(float(self.nrows)/self.batch_size)[0] > 0.0 :
            self.batch_size = int(float(self.nrows)/20)

        # Use 10% of the data each for testing and validating, the remaining for
        # training
        nfile_train = int(self.nfile*.8)
        nfile_test  = int(self.nfile*.1)
        nfile_val   = nfile_test

        n_data_check = self.nfile - ( nfile_train + nfile_test + nfile_val )
        if n_data_check > 0 :
            nfile_train += n_data_check
        elif n_data_check < 0 :
            nfile_train -= n_data_check

        TRAIN_DATA = np.zeros((nfile_train*self.nrows,self.ncols))
        train_images = np.zeros((nfile_train*self.nrows,self.ncols-1))
        train_labels = np.zeros((nfile_train*self.nrows,1))
        print 'Loading %d/%d files for training data...' % (nfile_train,self.nfile)
        for i in range(nfile_train) :
            print 'Loading file %d.' % (i+1)
            TRAIN_DATA[i*self.nrows:(i+1)*self.nrows,:] = np.loadtxt(self.full_file_path%(i+1))
        train_images = TRAIN_DATA[:,:-1].astype('int')
        train_labels = TRAIN_DATA[:,-1].astype('int')
        train_labels = convert_to_one_hot(train_labels)

        print 'Loading %d/%d files for test data...' % (nfile_test,self.nfile)
        TEST_DATA = np.zeros((nfile_test*self.nrows,self.ncols))
        test_images = np.zeros((nfile_test*self.nrows,self.ncols-1))
        test_labels = np.zeros((nfile_test*self.nrows,1))
        for i in range(nfile_test) :
            print 'Loading file %d.' % (i+1)
            TEST_DATA[i*self.nrows:(i+1)*self.nrows,:] = np.loadtxt(self.full_file_path%(i+1+nfile_train))
        test_images = TEST_DATA[:,:-1].astype('int')
        test_labels = TEST_DATA[:,-1].astype('int')
        if convert_test_labels_to_one_hot :
            test_labels = convert_to_one_hot(test_labels)

        print 'Loading %d/%d files for validation data...' % (nfile_val,self.nfile)
        VALIDATION_DATA = np.zeros((nfile_val*self.nrows,self.ncols))
        validation_images = np.zeros((nfile_val*self.nrows,self.ncols-1))
        validation_labels = np.zeros((nfile_val*self.nrows,1))
        for i in range(nfile_test) :
            print 'Loading file %d.' % (i+1)
            VALIDATION_DATA[i*self.nrows:(i+1)*self.nrows,:] = np.loadtxt(self.full_file_path%(i+1+nfile_train+nfile_test))
        validation_images = VALIDATION_DATA[:,:-1].astype('int')
        validation_labels = VALIDATION_DATA[:,-1].astype('int')

        data_sets.train = insert_file_info.DataSet(train_images, train_labels)
        data_sets.test = insert_file_info.DataSet(test_images, test_labels)
        data_sets.validation = insert_file_info.DataSet(validation_images, validation_labels)

        return data_sets

    def categorize_dose_of_data(self) :
        class DataSets(object):
            pass
        data_sets = DataSets()

        train_images = np.array([]).astype('int')
        train_labels = np.array([]).astype('int')

        test_images = np.array([]).astype('int')
        test_labels = np.array([]).astype('int')

        validation_images = np.array([]).astype('int')
        validation_labels = np.array([]).astype('int')

        data_sets.train = insert_file_info.DataSet(train_images, train_labels, data_type = 'train')
        data_sets.test = insert_file_info.DataSet(test_images, test_labels, data_type = 'test')
        data_sets.validation = insert_file_info.DataSet(validation_images, validation_labels, data_type = 'validation')

        return data_sets
<commit_msg>Add the ability to keep track of time when reading file.<commit_after>import numpy as np
import random
import sys

class insert_file_info :

    def __init__(self, full_file_path, filenumber, batch_size = 50,
        use_random_seed = False ) :
        """ full_file_path : full file path of the shuffled data
            filenumber     : An array of file number """
        self.filename         = full_file_path.rsplit('\\', 1)[-1]
        self.filename         = self.filename.rsplit('/', 1)[-1]
        self.full_file_path   = full_file_path
        self.nrows            = 0
        self.ncols            = 0
        self.nfile            = len(filenumber)
        self.batch_size       = batch_size
        self.current_index    = 0

    class DataSet(object) :
        file_info = None

        def __init__(self, images, labels, data_type = 'unknown') :
            #self.file_into = insert_file_info()

            #super(DataSet,self).__init__()
            #self.insert_file_info = insert_file_info

            self.data_type = data_type
            self.batch_size = 0
            self._ndata     = 0
            self._images = images
            self._labels = labels
            self._epochs_completed = 0
            self._index_in_epoch = 0
            self._index_in_datafile = 0
            self._file_index = 1

        #@staticmethod
        #def feed_self(self, batch_size, nrows) :
        #    self.batch_size = batch_size
        #    self.nrows      = nrows
            #print self.batch_size, self.nrows

        @property
        def images(self):
            return self._images

        @property
        def labels(self):
            return self._labels

        @property
        def ndata(self):
            return self._ndata

        @property
        def epochs_completed(self):
            return self._epochs_completed

        def next_batch(self, batch_size = 50, ndata = 8000) :
            self.batch_size = batch_size
            self._ndata     = ndata
            start = self._index_in_epoch
            if ( self._epochs_completed == 0 ) and ( start == 0 ) :
                print 'batch size : %d'    % self.batch_size
                print 'number of data: %d' % self._ndata
            self._index_in_epoch += self.batch_size
            if self._index_in_epoch > self._ndata :
                # Number of training epochs completed
                self._epochs_completed += 1
                # Shuffle data
                random.shuffle(self._images)
                random.shuffle(self._labels)
                # Reinitialize conunter
                start = 0
                self._index_in_epoch = self.batch_size
                assert self.batch_size <= self._ndata
            end = self._index_in_epoch
            return self._images[start:end], self._labels[start:end]

        def initialize_file_info_for_dose_of_data(self, full_file_path,
            filenumber, convert_test_labels_to_one_hot = True) :
            """ full_file_path : full file path of the shuffled data
                filenumber     : An array of file number """

            print 'Initializing file information for loading small dosage of data...'
            self.filename         = full_file_path.rsplit('\\', 1)[-1]
            self.filename         = self.filename.rsplit('/', 1)[-1]
            self.full_file_path   = full_file_path
            self.nfile            = len(filenumber)

            data = np.loadtxt(self.full_file_path%1)
            self.nrows, self.ncols = np.shape(data)
            self.nrows, self.ncols = int(self.nrows), int(self.ncols)

            # Use 10% of the data each for testing and validating, the remaining for
            # training
            self._nfile_train = int(self.nfile*.8)
            self._nfile_test  = int(self.nfile*.1)
            self._nfile_val   = self._nfile_test

            n_data_check = self.nfile - ( self._nfile_train + self._nfile_test + self._nfile_val )
            if n_data_check > 0 :
                self._nfile_train += n_data_check
            elif n_data_check < 0 :
                self._nfile_train -= n_data_check

            self._ndata           = self._nfile_train*self.nrows

            if self.data_type == 'train' :
                self.start_file_index = 1
                self.end_file_index = self._nfile_train
                self.convert_to_one_hot = True
            elif self.data_type == 'test' :
                self.start_file_index = self._nfile_train + 1
                self.end_file_index = self._nfile_train + self._nfile_test
                if convert_test_labels_to_one_hot :
                    self.convert_to_one_hot = convert_test_labels_to_one_hot
            elif self.data_type == 'validation' :
                self.start_file_index = self._nfile_train + self._nfile_test + 1
                self.end_file_index = self._nfile_train + self._nfile_test + self._nfile_validation
                self.convert_to_one_hot = False

        def next_dose(self, batch_size = 50) :

            def convert_to_one_hot( label ) :
                label_one_hot = np.zeros((len(label),2))
                for i in range(len(label)) :
                    label_one_hot[i,label[i]] = 1
                return label_one_hot

            start = self._index_in_datafile

            if ( self._file_index == self.start_file_index ) and ( start == 0 ) :
                self.batch_size = batch_size
                if np.modf(float(self.nrows)/self.batch_size)[0] > 0.0 :
                    self.batch_size = int(float(self.nrows)/20)

                print 'batch size : %d'    % self.batch_size
                print 'number of data: %d' % self._ndata

            self._index_in_datafile += self.batch_size
            if self._index_in_datafile > self.nrows :
                self._file_index += 1
                start = 0
                self._index_in_datafile = self.batch_size
                assert self.batch_size <= self.nrows

            if self._file_index > self.end_file_index :
                # Number of training epochs completed
                self._epochs_completed += 1
                self._file_index = self.start_file_index
                # Shufle data
                random.shuffle(self._images)
                random.shuffle(self._labels)
                # Reinitialize conunter
                start = 0
                self._index_in_datafile = self.batch_size

            end = self._index_in_datafile

            # Read in small dosage of data
            data = np.genfromtxt(self.full_file_path%(self._file_index) ,dtype=int,
                   skip_header=start, skip_footer=self.nrows-end)
            self._images = data[:,:-1].astype('int')
            labels = data[:,-1:].astype('int')
            if self.convert_to_one_hot :
                self._labels = convert_to_one_hot(labels)

            return self._images, self._labels

    def categorize_data(self, convert_test_labels_to_one_hot = True) :
        class DataSets(object):
            pass
        data_sets = DataSets()

        def convert_to_one_hot( label ) :
            label_one_hot = np.zeros((len(label),2))
            for i in range(len(label)) :
                label_one_hot[i,label[i]] = 1
            return label_one_hot

        data = np.loadtxt(self.full_file_path%1)
        self.nrows, self.ncols = np.shape(data)
        self.nrows, self.ncols = int(self.nrows), int(self.ncols)

        if np.modf(float(self.nrows)/self.batch_size)[0] > 0.0 :
            self.batch_size = int(float(self.nrows)/20)

        # Use 10% of the data each for testing and validating, the remaining for
        # training
        nfile_train = int(self.nfile*.8)
        nfile_test  = int(self.nfile*.1)
        nfile_val   = nfile_test

        n_data_check = self.nfile - ( nfile_train + nfile_test + nfile_val )
        if n_data_check > 0 :
            nfile_train += n_data_check
        elif n_data_check < 0 :
            nfile_train -= n_data_check

        start_time = time.time()

        TRAIN_DATA = np.zeros((nfile_train*self.nrows,self.ncols))
        train_images = np.zeros((nfile_train*self.nrows,self.ncols-1))
        train_labels = np.zeros((nfile_train*self.nrows,1))
        print 'Loading %d/%d files for training data...' % (nfile_train,self.nfile)
        for i in range(nfile_train) :
            print '%.1fs. Loading file %d.' % (time.time()-start_time, i+1)
            TRAIN_DATA[i*self.nrows:(i+1)*self.nrows,:] = np.loadtxt(self.full_file_path%(i+1))
        train_images = TRAIN_DATA[:,:-1].astype('int')
        train_labels = TRAIN_DATA[:,-1].astype('int')
        train_labels = convert_to_one_hot(train_labels)

        print 'Loading %d/%d files for test data...' % (nfile_test,self.nfile)
        TEST_DATA = np.zeros((nfile_test*self.nrows,self.ncols))
        test_images = np.zeros((nfile_test*self.nrows,self.ncols-1))
        test_labels = np.zeros((nfile_test*self.nrows,1))
        for i in range(nfile_test) :
            '%.1fs. Loading file %d.' % (time.time()-start_time, i+1)
            TEST_DATA[i*self.nrows:(i+1)*self.nrows,:] = np.loadtxt(self.full_file_path%(i+1+nfile_train))
        test_images = TEST_DATA[:,:-1].astype('int')
        test_labels = TEST_DATA[:,-1].astype('int')
        if convert_test_labels_to_one_hot :
            test_labels = convert_to_one_hot(test_labels)

        print 'Loading %d/%d files for validation data...' % (nfile_val,self.nfile)
        VALIDATION_DATA = np.zeros((nfile_val*self.nrows,self.ncols))
        validation_images = np.zeros((nfile_val*self.nrows,self.ncols-1))
        validation_labels = np.zeros((nfile_val*self.nrows,1))
        for i in range(nfile_test) :
            '%.1fs. Loading file %d.' % (time.time()-start_time, i+1)
            VALIDATION_DATA[i*self.nrows:(i+1)*self.nrows,:] = np.loadtxt(self.full_file_path%(i+1+nfile_train+nfile_test))
        validation_images = VALIDATION_DATA[:,:-1].astype('int')
        validation_labels = VALIDATION_DATA[:,-1].astype('int')

        data_sets.train = insert_file_info.DataSet(train_images, train_labels)
        data_sets.test = insert_file_info.DataSet(test_images, test_labels)
        data_sets.validation = insert_file_info.DataSet(validation_images, validation_labels)

        return data_sets

    def categorize_dose_of_data(self) :
        class DataSets(object):
            pass
        data_sets = DataSets()

        train_images = np.array([]).astype('int')
        train_labels = np.array([]).astype('int')

        test_images = np.array([]).astype('int')
        test_labels = np.array([]).astype('int')

        validation_images = np.array([]).astype('int')
        validation_labels = np.array([]).astype('int')

        data_sets.train = insert_file_info.DataSet(train_images, train_labels, data_type = 'train')
        data_sets.test = insert_file_info.DataSet(test_images, test_labels, data_type = 'test')
        data_sets.validation = insert_file_info.DataSet(validation_images, validation_labels, data_type = 'validation')

        return data_sets
<|endoftext|>


<commit_before>import oyaml as yaml
import sys
import configargparse

parser = configargparse.ArgumentParser(auto_env_var_prefix="INIT_")
parser.add_argument("--config-path", help="path to the configuration.yml file", default="/opt/opencga/conf/configuration.yml")
parser.add_argument("--client-config-path", help="path to the client-configuration.yml file", default="/opt/opencga/conf/client-configuration.yml")
parser.add_argument("--storage-config-path", help="path to the storage-configuration.yml file", default="/opt/opencga/conf/storage-configuration.yml")
parser.add_argument("--database-prefix", required=False)
parser.add_argument("--search-hosts", required=True)
parser.add_argument("--clinical-hosts", required=True)
parser.add_argument("--rga-hosts", required=False)
parser.add_argument("--cellbase-rest-url", required=False, help="Cellbase rest server hosting the cellbase service")
parser.add_argument("--cellbase-version", required=False, help="Cellbase rest server version (v4, v5..)")
parser.add_argument("--catalog-database-hosts", required=True)
parser.add_argument("--catalog-database-user", required=True)
parser.add_argument("--catalog-database-password", required=True)
parser.add_argument("--catalog-database-ssl", required=False, default=True)
parser.add_argument("--catalog-database-authentication-database", required=False, default="admin")
parser.add_argument("--catalog-database-authentication-mechanism", required=False)
parser.add_argument("--catalog-database-replica-set", required=False)
parser.add_argument("--catalog-search-hosts", required=True)
parser.add_argument("--catalog-search-user", required=False)
parser.add_argument("--catalog-search-password", required=False)
parser.add_argument("--rest-host", required=True)
parser.add_argument("--grpc-host", required=True)
parser.add_argument("--analysis-execution-mode", required=False)
parser.add_argument("--batch-account-name", required=False)
parser.add_argument("--batch-account-key", required=False)
parser.add_argument("--batch-endpoint", required=False)
parser.add_argument("--batch-pool-id", required=False)
parser.add_argument("--k8s-master-node", required=False)
parser.add_argument("--k8s-namespace", required=False, default="default")
parser.add_argument("--k8s-image-name", required=False)
parser.add_argument("--k8s-volumes-pvc-conf", required=False)
parser.add_argument("--k8s-volumes-pvc-sessions", required=False)
parser.add_argument("--k8s-volumes-pvc-variants", required=False)
parser.add_argument("--k8s-volumes-pvc-analysisconf", required=False)
parser.add_argument("--max-concurrent-jobs", required=False)
parser.add_argument("--variant-default-engine", required=False, default="hadoop")
parser.add_argument("--variant-options", required=False, action="append", nargs='*')
parser.add_argument("--hadoop-ssh-dns", required=False)
parser.add_argument("--hadoop-ssh-user", required=False)
parser.add_argument("--hadoop-ssh-pass", required=False)
parser.add_argument("--hadoop-ssh-key", required=False)
parser.add_argument("--hadoop-ssh-remote-opencga-home", required=False)
parser.add_argument("--health-check-interval", required=False)
parser.add_argument("--save", help="save update to source configuration files (default: false)", default=False, action='store_true')
args = parser.parse_args()

# TODO: Add check for a job config.

##############################################################################################################
# Load storage configuration yaml
##############################################################################################################

with open(args.storage_config_path) as f:
    storage_config = yaml.safe_load(f)

def hostOverride(hosts_var,server,component):
    hosts = hosts_var.replace('\"','').replace('[','').replace(']','').split(",")
    for i, host in enumerate(hosts):
        if component == "":
            if i == 0:
                storage_config[server]["hosts"].clear()
            storage_config[server]["hosts"].insert(i, hosts.strip())
        else:
            if i == 0:
                config[server][component]["hosts"].clear()
            config[server][component]["hosts"].insert(i, hosts.strip())

#  Inject search/clinical/rga hosts
hostOverride(args.search_hosts,"search","")
hostOverride(args.clinical_hosts,"clinical","")
hostOverride(args.rga_hosts,"rga","")

# Inject cellbase rest host, if set
if args.cellbase_rest_url is not None and args.cellbase_rest_url != "":
    cellbase_rest_url = args.cellbase_rest_url.replace('\"', '').replace('[','').replace(']','')
    storage_config["cellbase"]["url"] = cellbase_rest_url
if args.cellbase_version is not None and args.cellbase_version != "":
    storage_config["cellbase"]["version"] = args.cellbase_version

# set default engine
storage_config["variant"]["defaultEngine"] = args.variant_default_engine
storage_config["variant"]["options"]["annotator"] = "cellbase"

if args.variant_options is not None:
    # print("variant options: ", args.variant_options)
    for options in args.variant_options:
        for option in options:
            # print("option: ", option)
            kv = option.split("=")
            storage_config["variant"]["options"][kv[0]] = kv[1]

# Inject Hadoop ssh configuration
for _, storage_engine in enumerate(storage_config["variant"]["engines"]):

    if storage_engine["id"] == "hadoop":
        storage_engine["options"]["storage.hadoop.mr.executor"] = "ssh"
        storage_engine["options"]["storage.hadoop.mr.executor.ssh.host"] = args.hadoop_ssh_dns
        storage_engine["options"]["storage.hadoop.mr.executor.ssh.user"] = args.hadoop_ssh_user
        storage_engine["options"]["storage.hadoop.mr.executor.ssh.password"] = args.hadoop_ssh_pass
        if  args.hadoop_ssh_key is not None and  args.hadoop_ssh_key != "":
            storage_engine["options"]["storage.hadoop.mr.executor.ssh.key"] = args.hadoop_ssh_key
        else:
            storage_engine["options"]["storage.hadoop.mr.executor.ssh.key"] = ""

        storage_engine["options"]["storage.hadoop.mr.executor.ssh.remoteOpenCgaHome"] = args.hadoop_ssh_remote_opencga_home

##############################################################################################################
# Load configuration yaml
##############################################################################################################

with open(args.config_path) as f:
    config = yaml.safe_load(f)

if args.database_prefix:
    config["databasePrefix"] = args.database_prefix

# Inject catalog database
hostOverride(args.catalog_database_hosts,"catalog","database")

config["catalog"]["database"]["user"] = args.catalog_database_user
config["catalog"]["database"]["password"] = args.catalog_database_password
config["catalog"]["database"]["options"]["sslEnabled"] = args.catalog_database_ssl
config["catalog"]["database"]["options"]["sslInvalidCertificatesAllowed"] = True
if args.catalog_database_replica_set is not None and args.catalog_database_replica_set != "":
    config["catalog"]["database"]["options"]["replicaSet"] = args.catalog_database_replica_set
if args.catalog_database_authentication_database is not None:
    config["catalog"]["database"]["options"]["authenticationDatabase"] = args.catalog_database_authentication_database
if args.catalog_database_authentication_mechanism is not None:
    config["catalog"]["database"]["options"]["authenticationMechanism"] = args.catalog_database_authentication_mechanism

# Inject search database
hostOverride(args.catalog_search_hosts,"catalog","searchEngine")

if args.catalog_search_user is not None:
    config["catalog"]["searchEngine"]["user"] = args.catalog_search_user
    config["catalog"]["searchEngine"]["password"] = args.catalog_search_password

# Inject execution settings
config["analysis"]["scratchDir"] = "/tmp/opencga_scratch"
if args.max_concurrent_jobs is not None:
    config["analysis"]["execution"]["maxConcurrentJobs"]["variant-index"] = int(args.max_concurrent_jobs)

if args.analysis_execution_mode is not None:
    config["analysis"]["execution"]["id"] = args.analysis_execution_mode

if args.analysis_execution_mode == "AZURE":
    config["analysis"]["execution"]["options"] = {}
    config["analysis"]["execution"]["options"]["azure.batchAccount"] = args.batch_account_name
    config["analysis"]["execution"]["options"]["azure.batchKey"] = args.batch_account_key
    config["analysis"]["execution"]["options"]["azure.batchUri"] = args.batch_endpoint
    config["analysis"]["execution"]["options"]["azure.batchPoolId"] = args.batch_pool_id
elif args.analysis_execution_mode == "k8s":
    config["analysis"]["execution"]["options"]["k8s.masterUrl"] = args.k8s_master_node
    config["analysis"]["execution"]["options"]["k8s.namespace"] = args.k8s_namespace

    if args.k8s_image_name:
        config["analysis"]["execution"]["options"]["k8s.imageName"] = args.k8s_image_name

    config["analysis"]["execution"]["options"]["k8s.volumes"][0]["name"] = "conf"
    config["analysis"]["execution"]["options"]["k8s.volumes"][0]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_conf
    config["analysis"]["execution"]["options"]["k8s.volumes"][1]["name"] = "sessions"
    config["analysis"]["execution"]["options"]["k8s.volumes"][1]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_sessions
    config["analysis"]["execution"]["options"]["k8s.volumes"][2]["name"] = "variants"
    config["analysis"]["execution"]["options"]["k8s.volumes"][2]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_variants
    config["analysis"]["execution"]["options"]["k8s.volumes"][3]["name"] = "analysisconf"
    config["analysis"]["execution"]["options"]["k8s.volumes"][3]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_analysisconf

# Inject healthCheck interval
if args.health_check_interval is not None:
    config["healthCheck"]["interval"] = args.health_check_interval

##############################################################################################################
# Load client configuration yaml
##############################################################################################################

with open(args.client_config_path) as f:
    client_config = yaml.safe_load(f)

# Inject grpc and rest host
client_config["rest"]["hosts"] = [{}]
client_config["rest"]["hosts"][0]["name"] = "default"
client_config["rest"]["hosts"][0]["url"] = args.rest_host.replace('"','')
client_config["rest"]["defaultHostIndex"] = 0
client_config["grpc"]["host"] = args.grpc_host.replace('"','')

# Running with --save will update the configuration files inplace.
# Without --save will simply dump the update YAML to stdout so that
# the caller can handle it.
# Note: The dump will use the safe representation so there is likely
# to be format diffs between the original input and the output as well
# as value changes.
if args.save == False:
    yaml.dump(storage_config, sys.stdout, default_flow_style=False, allow_unicode=True)
    print("---")  # Add yaml delimiter
    yaml.dump(config, sys.stdout, default_flow_style=False, allow_unicode=True)
    print("---")  # Add yaml delimiter
    yaml.dump(client_config, sys.stdout, default_flow_style=False, allow_unicode=True)
else:
    with open(args.storage_config_path, "w") as f:
        yaml.dump(storage_config, f, default_flow_style=False)
    with open(args.config_path, "w") as f:
        yaml.dump(config, f, default_flow_style=False)
    with open(args.client_config_path, "w") as f:
        yaml.dump(client_config, f, default_flow_style=False)

<commit_msg>Update override_yaml add method hostOverride #TASK-923<commit_after>import oyaml as yaml
import sys
import configargparse

parser = configargparse.ArgumentParser(auto_env_var_prefix="INIT_")
parser.add_argument("--config-path", help="path to the configuration.yml file", default="/opt/opencga/conf/configuration.yml")
parser.add_argument("--client-config-path", help="path to the client-configuration.yml file", default="/opt/opencga/conf/client-configuration.yml")
parser.add_argument("--storage-config-path", help="path to the storage-configuration.yml file", default="/opt/opencga/conf/storage-configuration.yml")
parser.add_argument("--database-prefix", required=False)
parser.add_argument("--search-hosts", required=True)
parser.add_argument("--clinical-hosts", required=True)
parser.add_argument("--rga-hosts", required=False)
parser.add_argument("--cellbase-rest-url", required=False, help="Cellbase rest server hosting the cellbase service")
parser.add_argument("--cellbase-version", required=False, help="Cellbase rest server version (v4, v5..)")
parser.add_argument("--catalog-database-hosts", required=True)
parser.add_argument("--catalog-database-user", required=True)
parser.add_argument("--catalog-database-password", required=True)
parser.add_argument("--catalog-database-ssl", required=False, default=True)
parser.add_argument("--catalog-database-authentication-database", required=False, default="admin")
parser.add_argument("--catalog-database-authentication-mechanism", required=False)
parser.add_argument("--catalog-database-replica-set", required=False)
parser.add_argument("--catalog-search-hosts", required=True)
parser.add_argument("--catalog-search-user", required=False)
parser.add_argument("--catalog-search-password", required=False)
parser.add_argument("--rest-host", required=True)
parser.add_argument("--grpc-host", required=True)
parser.add_argument("--analysis-execution-mode", required=False)
parser.add_argument("--batch-account-name", required=False)
parser.add_argument("--batch-account-key", required=False)
parser.add_argument("--batch-endpoint", required=False)
parser.add_argument("--batch-pool-id", required=False)
parser.add_argument("--k8s-master-node", required=False)
parser.add_argument("--k8s-namespace", required=False, default="default")
parser.add_argument("--k8s-image-name", required=False)
parser.add_argument("--k8s-volumes-pvc-conf", required=False)
parser.add_argument("--k8s-volumes-pvc-sessions", required=False)
parser.add_argument("--k8s-volumes-pvc-variants", required=False)
parser.add_argument("--k8s-volumes-pvc-analysisconf", required=False)
parser.add_argument("--max-concurrent-jobs", required=False)
parser.add_argument("--variant-default-engine", required=False, default="hadoop")
parser.add_argument("--variant-options", required=False, action="append", nargs='*')
parser.add_argument("--hadoop-ssh-dns", required=False)
parser.add_argument("--hadoop-ssh-user", required=False)
parser.add_argument("--hadoop-ssh-pass", required=False)
parser.add_argument("--hadoop-ssh-key", required=False)
parser.add_argument("--hadoop-ssh-remote-opencga-home", required=False)
parser.add_argument("--health-check-interval", required=False)
parser.add_argument("--save", help="save update to source configuration files (default: false)", default=False, action='store_true')
args = parser.parse_args()

# TODO: Add check for a job config.

##############################################################################################################
# Load storage configuration yaml
##############################################################################################################

with open(args.storage_config_path) as f:
    storage_config = yaml.safe_load(f)

def hostOverride(hosts_var,server,component):
    hosts = hosts_var.replace('\"','').replace('[','').replace(']','').split(",")
    for i, host in enumerate(hosts):
        if component == "":
            if i == 0:
                storage_config[server]["hosts"].clear()
            storage_config[server]["hosts"].insert(i, host.strip())
        else:
            if i == 0:
                config[server][component]["hosts"].clear()
            config[server][component]["hosts"].insert(i, host.strip())

#  Inject search/clinical/rga hosts
hostOverride(args.search_hosts,"search","")
hostOverride(args.clinical_hosts,"clinical","")
hostOverride(args.rga_hosts,"rga","")

# Inject cellbase rest host, if set
if args.cellbase_rest_url is not None and args.cellbase_rest_url != "":
    cellbase_rest_url = args.cellbase_rest_url.replace('\"', '').replace('[','').replace(']','')
    storage_config["cellbase"]["url"] = cellbase_rest_url
if args.cellbase_version is not None and args.cellbase_version != "":
    storage_config["cellbase"]["version"] = args.cellbase_version

# set default engine
storage_config["variant"]["defaultEngine"] = args.variant_default_engine
storage_config["variant"]["options"]["annotator"] = "cellbase"

if args.variant_options is not None:
    # print("variant options: ", args.variant_options)
    for options in args.variant_options:
        for option in options:
            # print("option: ", option)
            kv = option.split("=")
            storage_config["variant"]["options"][kv[0]] = kv[1]

# Inject Hadoop ssh configuration
for _, storage_engine in enumerate(storage_config["variant"]["engines"]):

    if storage_engine["id"] == "hadoop":
        storage_engine["options"]["storage.hadoop.mr.executor"] = "ssh"
        storage_engine["options"]["storage.hadoop.mr.executor.ssh.host"] = args.hadoop_ssh_dns
        storage_engine["options"]["storage.hadoop.mr.executor.ssh.user"] = args.hadoop_ssh_user
        storage_engine["options"]["storage.hadoop.mr.executor.ssh.password"] = args.hadoop_ssh_pass
        if  args.hadoop_ssh_key is not None and  args.hadoop_ssh_key != "":
            storage_engine["options"]["storage.hadoop.mr.executor.ssh.key"] = args.hadoop_ssh_key
        else:
            storage_engine["options"]["storage.hadoop.mr.executor.ssh.key"] = ""

        storage_engine["options"]["storage.hadoop.mr.executor.ssh.remoteOpenCgaHome"] = args.hadoop_ssh_remote_opencga_home

##############################################################################################################
# Load configuration yaml
##############################################################################################################

with open(args.config_path) as f:
    config = yaml.safe_load(f)

if args.database_prefix:
    config["databasePrefix"] = args.database_prefix

# Inject catalog database
hostOverride(args.catalog_database_hosts,"catalog","database")

config["catalog"]["database"]["user"] = args.catalog_database_user
config["catalog"]["database"]["password"] = args.catalog_database_password
config["catalog"]["database"]["options"]["sslEnabled"] = args.catalog_database_ssl
config["catalog"]["database"]["options"]["sslInvalidCertificatesAllowed"] = True
if args.catalog_database_replica_set is not None and args.catalog_database_replica_set != "":
    config["catalog"]["database"]["options"]["replicaSet"] = args.catalog_database_replica_set
if args.catalog_database_authentication_database is not None:
    config["catalog"]["database"]["options"]["authenticationDatabase"] = args.catalog_database_authentication_database
if args.catalog_database_authentication_mechanism is not None:
    config["catalog"]["database"]["options"]["authenticationMechanism"] = args.catalog_database_authentication_mechanism

# Inject search database
hostOverride(args.catalog_search_hosts,"catalog","searchEngine")

if args.catalog_search_user is not None:
    config["catalog"]["searchEngine"]["user"] = args.catalog_search_user
    config["catalog"]["searchEngine"]["password"] = args.catalog_search_password

# Inject execution settings
config["analysis"]["scratchDir"] = "/tmp/opencga_scratch"
if args.max_concurrent_jobs is not None:
    config["analysis"]["execution"]["maxConcurrentJobs"]["variant-index"] = int(args.max_concurrent_jobs)

if args.analysis_execution_mode is not None:
    config["analysis"]["execution"]["id"] = args.analysis_execution_mode

if args.analysis_execution_mode == "AZURE":
    config["analysis"]["execution"]["options"] = {}
    config["analysis"]["execution"]["options"]["azure.batchAccount"] = args.batch_account_name
    config["analysis"]["execution"]["options"]["azure.batchKey"] = args.batch_account_key
    config["analysis"]["execution"]["options"]["azure.batchUri"] = args.batch_endpoint
    config["analysis"]["execution"]["options"]["azure.batchPoolId"] = args.batch_pool_id
elif args.analysis_execution_mode == "k8s":
    config["analysis"]["execution"]["options"]["k8s.masterUrl"] = args.k8s_master_node
    config["analysis"]["execution"]["options"]["k8s.namespace"] = args.k8s_namespace

    if args.k8s_image_name:
        config["analysis"]["execution"]["options"]["k8s.imageName"] = args.k8s_image_name

    config["analysis"]["execution"]["options"]["k8s.volumes"][0]["name"] = "conf"
    config["analysis"]["execution"]["options"]["k8s.volumes"][0]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_conf
    config["analysis"]["execution"]["options"]["k8s.volumes"][1]["name"] = "sessions"
    config["analysis"]["execution"]["options"]["k8s.volumes"][1]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_sessions
    config["analysis"]["execution"]["options"]["k8s.volumes"][2]["name"] = "variants"
    config["analysis"]["execution"]["options"]["k8s.volumes"][2]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_variants
    config["analysis"]["execution"]["options"]["k8s.volumes"][3]["name"] = "analysisconf"
    config["analysis"]["execution"]["options"]["k8s.volumes"][3]["persistentVolumeClaim"]["claimName"] = args.k8s_volumes_pvc_analysisconf

# Inject healthCheck interval
if args.health_check_interval is not None:
    config["healthCheck"]["interval"] = args.health_check_interval

##############################################################################################################
# Load client configuration yaml
##############################################################################################################

with open(args.client_config_path) as f:
    client_config = yaml.safe_load(f)

# Inject grpc and rest host
client_config["rest"]["hosts"] = [{}]
client_config["rest"]["hosts"][0]["name"] = "default"
client_config["rest"]["hosts"][0]["url"] = args.rest_host.replace('"','')
client_config["rest"]["defaultHostIndex"] = 0
client_config["grpc"]["host"] = args.grpc_host.replace('"','')

# Running with --save will update the configuration files inplace.
# Without --save will simply dump the update YAML to stdout so that
# the caller can handle it.
# Note: The dump will use the safe representation so there is likely
# to be format diffs between the original input and the output as well
# as value changes.
if args.save == False:
    yaml.dump(storage_config, sys.stdout, default_flow_style=False, allow_unicode=True)
    print("---")  # Add yaml delimiter
    yaml.dump(config, sys.stdout, default_flow_style=False, allow_unicode=True)
    print("---")  # Add yaml delimiter
    yaml.dump(client_config, sys.stdout, default_flow_style=False, allow_unicode=True)
else:
    with open(args.storage_config_path, "w") as f:
        yaml.dump(storage_config, f, default_flow_style=False)
    with open(args.config_path, "w") as f:
        yaml.dump(config, f, default_flow_style=False)
    with open(args.client_config_path, "w") as f:
        yaml.dump(client_config, f, default_flow_style=False)

<|endoftext|>


<commit_before>from django.core.management.base import BaseCommand
from corehq import Domain

class Command(BaseCommand):
    help = 'Adds domains that have apps w/ media to that media files list of valid domains'

    def handle(self, *args, **options):
        for d in Domain.get_all():
            for app in d.full_applications():
                if app.is_remote_app():
                    continue
                for _, m in app.get_media_objects():
                    if app.domain not in m.valid_domains:
                        m.valid_domains.append(app.domain)
                        self.stdout.write("adding domain %s to media file %s" % (app.domain, m._id))
                        m.save()
<commit_msg>more robust set_media_domains script<commit_after>from optparse import make_option
from django.core.management.base import BaseCommand
from corehq import Domain

class Command(BaseCommand):
    help = 'Adds domains that have apps w/ media to that media files list of valid domains'

    option_list = BaseCommand.option_list + (
        make_option('-s', '--since',
            help="Begin at this domain. (Alphabetically Descending"),
        make_option('-d', '--domain',
            help="Check only this domain"),
        )

    def handle(self, *args, **options):
        if options.get('domain'):
            domains = [Domain.get_by_name(options.get('domain'))]
        else:
            domains = Domain.get_all()

        since = options.get('since')
        seen_since = False
        for d in domains:
            if since and not seen_since:
                if d.name == since:
                    seen_since = True
                else:
                    continue

            for app in d.full_applications():
                try:
                    if app.is_remote_app():
                        continue
                    for _, m in app.get_media_objects():
                        if app.domain not in m.valid_domains:
                            m.valid_domains.append(app.domain)
                            self.stdout.write("adding domain %s to media file %s" % (app.domain, m._id))
                            m.save()
                except Exception as e:
                    self.stdout.write("Error in app %s-%s: %s" % (app.domain, app._id, e))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
#  init_stores.py
#  streamflow_prediction_tool
#
#  Created by Alan D. Snow, 2015-2017
#  License: BSD 3-Clause

from .app import StreamflowPredictionTool as app
from .model import Base, DataStore, DataStoreType

def init_main_db(engine, first_time):
    # Create tables
    Base.metadata.create_all(engine)

    # Initial data
    if first_time:
        # make session
        session_maker = \
            app.get_persistent_store_database('main_db',
                                              as_sessionmaker=True)
        session = session_maker()

        # add all possible data story types
        local_ds_type = \
            DataStoreType(code_name="local",
                          human_readable_name="Computations Local (None)")
        session.add(local_ds_type)
        session.add(DataStoreType(code_name="ckan",
                                  human_readable_name="CKAN"))
        session.add(DataStoreType(code_name="hydroshare",
                                  human_readable_name="HydroShare"))

        # add all possible data stores
        session.add(DataStore(name="Local Server",
                              owner_org="",
                              data_store_type=local_ds_type,
                              api_endpoint="local",
                              api_key=""))

        session.commit()
        session.close()
<commit_msg>remove hydroshare as unsupported<commit_after># -*- coding: utf-8 -*-
#
#  init_stores.py
#  streamflow_prediction_tool
#
#  Created by Alan D. Snow, 2015-2017
#  License: BSD 3-Clause

from .app import StreamflowPredictionTool as app
from .model import Base, DataStore, DataStoreType

def init_main_db(engine, first_time):
    # Create tables
    Base.metadata.create_all(engine)

    # Initial data
    if first_time:
        # make session
        session_maker = \
            app.get_persistent_store_database('main_db',
                                              as_sessionmaker=True)
        session = session_maker()

        # add all possible data story types
        local_ds_type = \
            DataStoreType(code_name="local",
                          human_readable_name="Computations Local (None)")
        session.add(local_ds_type)
        session.add(DataStoreType(code_name="ckan",
                                  human_readable_name="CKAN"))

        # add all possible data stores
        session.add(DataStore(name="Local Server",
                              owner_org="",
                              data_store_type=local_ds_type,
                              api_endpoint="local",
                              api_key=""))

        session.commit()
        session.close()
<|endoftext|>


<commit_before># http://blog.shayanjaved.com/2012/01/14/binary-search-tree-in-python/

# http://interactivepython.org/XSKWZ/LpOMZ/courselib/static/pythonds/Trees/bst.html
import timeit
import random
import subprocess
from collections import deque

class BSTNode(object):
    '''Instantiate Node and add helper functions'''
    def __init__(self, value, parent=None, left_child=None, right_child=None):
        self.value = value
        self.parent = parent
        self.left = left_child
        self.right = right_child
        self.height = 0
        self.balance = 0

    def is_root(self):
        '''Helper function for root node'''
        return not self.parent

    def is_leaf(self):
        ''''Helper function for acknowledging leaf'''
        return not (self.right_child or self.left_child)

    def is_left(self):
        '''Helper fuction for finding left child. Might be redundent with is_left
           Will decide later........'''
        if self.parent is None:
            return self.parent
        else:
            return self is self.parent.left_child

    def update_height(self, bubble_up=True):
        '''If bubble_up is True, we go up the tree correcting height/balance
           if not we will just correct the node'''
        if self.left_child is None:
            # set the left tree to zero
            left_height = 0
        else:
            left_height = self.left_child.height + 1
        if self.right_child is None:
            # set the right tree to zero
            right_height = 0
        else:
            right_height = self.right_child.height + 1
            # we want to be able to balance even if we don't change the height

        self.balance = left_height - right_height
        height = max(left_height, right_height)
        if self.height != height:
            self.height = height
            if self.parent is not None:
                # We only bubble up if the height changes
                if bubble_up:
                    self.parent.update_height()

    def _get_dot(self):
        """recursively prepare a dot graph entry for this node."""
        if self.left is not None:
            yield "\t%s -> %s;" % (self.value, self.left.value)
            for i in self.left._get_dot():
                yield i
        elif self.right is not None:
            r = random.randint(0, 1e9)
            yield "\tnull%s [shape=point];" % r
            yield "\t%s -> null%s;" % (self.value, r)
        if self.right is not None:
            yield "\t%s -> %s;" % (self.value, self.right.value)
            for i in self.right._get_dot():
                yield i
        elif self.left is not None:
            r = random.randint(0, 1e9)
            yield "\tnull%s [shape=point];" % r
            yield "\t%s -> null%s;" % (self.value, r)

class BST(object):
    '''Instantiate binary search tree'''
    def __init__(self, values=None):
        self.root = None
        self._size = 0

    def size(self):
        '''Will return integer size of BST'''
        return self._size

    def insert(self, value):
        '''Inserts the data in value into BST'''
        if self.root is None:
            self.root = BSTNode(value)
            self._size += 1
            return
        current_node = self.root
        while True:
            if current_node.value > value:
                if current_node.left:
                    current_node = current_node.left
                else:
                    current_node.left = BSTNode(value)
                    self._size += 1
                    break
            elif current_node.value < value:
                if current_node.right:
                    current_node = current_node.right
                else:
                    current_node.right = BSTNode(value)
                    self._size += 1
                    break
            else:
                break

    def contains(self, value):
        '''Returns true if data in value is in BST'''
        if self.root is None:
            return False
        current_node = self.root
        while True:
            if current_node.value > value:
                if current_node.left:
                    current_node = current_node.left
                else:
                    return False
            elif current_node.value < value:
                if current_node.right:
                    current_node = current_node.right
                else:
                    return False
            else:
                return True

    def depth(self):
        '''Returns total number of levels in BST as interger'''
        if self.root is None:
            return 0
        return self._depth(1, self.root)

    def _depth(self, curr_depth, local_root):
        '''Helper function for depth'''
        l_depth = r_depth = 0
        if local_root.left:
            l_depth = self._depth(curr_depth + 1, local_root.left)
        if local_root.right:
            r_depth = self._depth(curr_depth + 1, local_root.right)
        return max(curr_depth, l_depth, r_depth)

    def is_balanced(self):
        '''Return positive or negative integer to represent tree balance'''
        ret_value = 0
        if self.root is None:
            return ret_value
        if self.root.left:
            ret_value += self._depth(1, self.root.left)
        if self.root.right:
            ret_value -= self._depth(1, self.root.right)
        return ret_value

    def get_dot(self):
        """return the tree with root 'self' as a dot graph for visualization"""
        return "digraph G{\n%s}" % ("" if self.root.value is None else (
            "\t%s;\n%s\n" % (
                self.root.value,
                "\n".join(self.root._get_dot())
            )
        ))

    def in_order(self):
        return self._in_order(self.root)

    def _in_order(self, leaf):
        if leaf is None:
            return
        for value in self._in_order(leaf.left):
            yield value
        yield leaf.value
        for value in self._in_order(leaf.right):
            yield value

    def pre_order(self):
        return self._pre_order(self.root)

    def _pre_order(self, leaf):
        if leaf is None:
            return
        yield leaf.value
        for value in self._pre_order(leaf.left):
            yield value
        for value in self._pre_order(leaf.right):
            yield value

    def post_order(self):
        return self._post_order(self.root)

    def _post_order(self, leaf):
        if leaf is None:
            return
        for value in self._post_order(leaf.left):
            yield value
        for value in self._post_order(leaf.right):
            yield value
        yield leaf.value

    def breadth_traversal(self):
        x = deque()
        x.append(self.root)
        while x:
            leaf = x.popleft()
            yield leaf.value
            if leaf.left:
                x.append(leaf.left)
            if leaf.right:
                x.append(leaf.right)

    def delete(self, val):
        self.root = self._delete(val, self.root)
        return None

    def _delete(self, val, leaf):
        def _descendants(leaf):
            if leaf.left:
                return _descendants(leaf.left)
            else:
                return leaf.value

        if not leaf:
            return None

        if leaf.value == val:
            self._size -= 1
            if leaf.left and leaf.right:
                leaf.value = _descendants(leaf.right)
                leaf.right = self._delete(leaf.value, leaf.right)
                return leaf
            elif leaf.left and not leaf.right:
                return leaf.left
            elif not leaf.left and leaf.right:
                return leaf.right
            else:
                return None

        elif leaf.value < val:
            if leaf.right:
                leaf.right = self._delete(val, leaf.right)
            return leaf

        else:
            if leaf.left:
                leaf.left = self._delete(val, leaf.left)
            return leaf

    def rotate_left(self, root):
        left = root.is_left()
        pivot = root.right_child

        if pivot is None:
            return
        root.right_child = pivot.left_child
        if pivot.left_child is not None:
            root.right_child.parent = root
        pivot.left_child = root
        pivot.parent = root.parent
        root.parent = pivot
        if left is None:
            self.root = pivot
        elif left:
            pivot.parent.left_child = pivot
        else:
            pivot.parent.right_child = pivot
        root.update_height(False)
        pivot.update_height(False)

    def rotate_right(self, root):
        left = root.is_left()
        pivot = root.left_child
        if pivot is None:
            return
        root.left_child = pivot.right_child
        if pivot.right_child is not None:
            root.left_child.parent = root
        pivot.right_child = root
        pivot.parent = root.parent
        root.parent = pivot
        if left is None:
            self.root = pivot
        elif left:
            pivot.parent.left_child = pivot
        else:
            pivot.parent.right_child = pivot
        root.update_height(False)
        pivot.update_height(False)

    def find_leftmost(self, node):
        if node.left_child is None:
            return node
        else:
            return self.find_leftmost(node.left_child)

    def find_rightmost(self, node):
        if node.right_child is None:
            return node
        else:
            return self.find_rightmost(node.right_child)

    def find_next(self, value):
        node = self.find(value)
        if (node is None) or (node.value != value):
            return None
        else:
            right_child = node.right_child
            if right_child is not None:
                node = self.find_leftmost(right_child)
            else:
                parent = node.parent
                while(parent is not None):
                    if node is parent.left_child:
                        break
                    node = parent
                    parent = node.parent
                node = parent
            if node is None:
                return node
            else:
                return node.value

    def find_next(self, value):
        node = self.find(value)
        if (node is None) or (node.value != value):
            return None
        else:
            left_child = node.left_child
            if left_child is not None:
                node = self.find_leftmost(left_child)
            else:
                parent = node.parent
                while(parent is not None):
                    if node is parent.right_child:
                        break
                    node = parent
                    parent = node.parent
                node = parent
            if node is None:
                return node
            else:
                return node.value

    def balance(self, node):
        ''' There are four posabilities for rotation
            left-left=LL right-right=RR
            left-right=LR right-left=RL'''
        node.update_height(False)
        if node.balance == 2:
            if node.left_child.balance != -1:
                # LL rotation
                self.rotate_right(node)
                if node.parent.parent is not None:
                    self.balance(node.parent.parent)
            else:
                # LR rotation
                self.rotate_left(node.left_child)
                self.balance(node)
        elif node.balance == -2:
            if node.right_child.balance != 1:
                # RR rotation
                self.rotate_left(node)
                if node.parent.parent is not None:
                    self.balance(node.parent.parent)

            else:
                # RL rotation
                self.rotate_right(node.right_child)
                self.balance(node)
        else:
            if node.parent is not None:
                self.balance(node.parent)

    def sort(tree_maker, ascending=True):
        t = BST()
        for item in tree_maker:
            t.insert(item)
        ret = []
        if ascending:
            node = t.

if __name__ == '__main__':

    # x = range(100)
    # bst = BST()
    # for i in x:
    #     bst.insert(i)
    # bst.insert(42.1)
    # dot_graph = bst.get_dot()
    # t = subprocess.Popen(["dot", "-Tpng"], stdin=subprocess.PIPE)
    # t.communicate(dot_graph)

    def easy_tree():
        x = random.sample(range(100), 100)
        bst = BST()
        bst.insert(50)
        for i in x:
            bst.insert(i)
        bst.insert(42.1)
        bst.contains(42.1)

    def hard_tree():
        x = range(100)
        bst = BST()
        for i in x:
            bst.insert(i)
        bst.insert(42.1)
        bst.contains(42.1)

    print(timeit.Timer("easy_tree()", setup="from __main__ import easy_tree").timeit(number=1000))
    print(timeit.Timer("hard_tree()", setup="from __main__ import hard_tree").timeit(number=1000))
<commit_msg>corrected name for find_prev function<commit_after># http://blog.shayanjaved.com/2012/01/14/binary-search-tree-in-python/

# http://interactivepython.org/XSKWZ/LpOMZ/courselib/static/pythonds/Trees/bst.html
import timeit
import random
import subprocess
from collections import deque

class BSTNode(object):
    '''Instantiate Node and add helper functions'''
    def __init__(self, value, parent=None, left_child=None, right_child=None):
        self.value = value
        self.parent = parent
        self.left = left_child
        self.right = right_child
        self.height = 0
        self.balance = 0

    def is_root(self):
        '''Helper function for root node'''
        return not self.parent

    def is_leaf(self):
        ''''Helper function for acknowledging leaf'''
        return not (self.right_child or self.left_child)

    def is_left(self):
        '''Helper fuction for finding left child. Might be redundent with is_left
           Will decide later........'''
        if self.parent is None:
            return self.parent
        else:
            return self is self.parent.left_child

    def update_height(self, bubble_up=True):
        '''If bubble_up is True, we go up the tree correcting height/balance
           if not we will just correct the node'''
        if self.left_child is None:
            # set the left tree to zero
            left_height = 0
        else:
            left_height = self.left_child.height + 1
        if self.right_child is None:
            # set the right tree to zero
            right_height = 0
        else:
            right_height = self.right_child.height + 1
            # we want to be able to balance even if we don't change the height

        self.balance = left_height - right_height
        height = max(left_height, right_height)
        if self.height != height:
            self.height = height
            if self.parent is not None:
                # We only bubble up if the height changes
                if bubble_up:
                    self.parent.update_height()

    def _get_dot(self):
        """recursively prepare a dot graph entry for this node."""
        if self.left is not None:
            yield "\t%s -> %s;" % (self.value, self.left.value)
            for i in self.left._get_dot():
                yield i
        elif self.right is not None:
            r = random.randint(0, 1e9)
            yield "\tnull%s [shape=point];" % r
            yield "\t%s -> null%s;" % (self.value, r)
        if self.right is not None:
            yield "\t%s -> %s;" % (self.value, self.right.value)
            for i in self.right._get_dot():
                yield i
        elif self.left is not None:
            r = random.randint(0, 1e9)
            yield "\tnull%s [shape=point];" % r
            yield "\t%s -> null%s;" % (self.value, r)

class BST(object):
    '''Instantiate binary search tree'''
    def __init__(self, values=None):
        self.root = None
        self._size = 0

    def size(self):
        '''Will return integer size of BST'''
        return self._size

    def insert(self, value):
        '''Inserts the data in value into BST'''
        if self.root is None:
            self.root = BSTNode(value)
            self._size += 1
            return
        current_node = self.root
        while True:
            if current_node.value > value:
                if current_node.left:
                    current_node = current_node.left
                else:
                    current_node.left = BSTNode(value)
                    self._size += 1
                    break
            elif current_node.value < value:
                if current_node.right:
                    current_node = current_node.right
                else:
                    current_node.right = BSTNode(value)
                    self._size += 1
                    break
            else:
                break

    def contains(self, value):
        '''Returns true if data in value is in BST'''
        if self.root is None:
            return False
        current_node = self.root
        while True:
            if current_node.value > value:
                if current_node.left:
                    current_node = current_node.left
                else:
                    return False
            elif current_node.value < value:
                if current_node.right:
                    current_node = current_node.right
                else:
                    return False
            else:
                return True

    def depth(self):
        '''Returns total number of levels in BST as interger'''
        if self.root is None:
            return 0
        return self._depth(1, self.root)

    def _depth(self, curr_depth, local_root):
        '''Helper function for depth'''
        l_depth = r_depth = 0
        if local_root.left:
            l_depth = self._depth(curr_depth + 1, local_root.left)
        if local_root.right:
            r_depth = self._depth(curr_depth + 1, local_root.right)
        return max(curr_depth, l_depth, r_depth)

    def is_balanced(self):
        '''Return positive or negative integer to represent tree balance'''
        ret_value = 0
        if self.root is None:
            return ret_value
        if self.root.left:
            ret_value += self._depth(1, self.root.left)
        if self.root.right:
            ret_value -= self._depth(1, self.root.right)
        return ret_value

    def get_dot(self):
        """return the tree with root 'self' as a dot graph for visualization"""
        return "digraph G{\n%s}" % ("" if self.root.value is None else (
            "\t%s;\n%s\n" % (
                self.root.value,
                "\n".join(self.root._get_dot())
            )
        ))

    def in_order(self):
        return self._in_order(self.root)

    def _in_order(self, leaf):
        if leaf is None:
            return
        for value in self._in_order(leaf.left):
            yield value
        yield leaf.value
        for value in self._in_order(leaf.right):
            yield value

    def pre_order(self):
        return self._pre_order(self.root)

    def _pre_order(self, leaf):
        if leaf is None:
            return
        yield leaf.value
        for value in self._pre_order(leaf.left):
            yield value
        for value in self._pre_order(leaf.right):
            yield value

    def post_order(self):
        return self._post_order(self.root)

    def _post_order(self, leaf):
        if leaf is None:
            return
        for value in self._post_order(leaf.left):
            yield value
        for value in self._post_order(leaf.right):
            yield value
        yield leaf.value

    def breadth_traversal(self):
        x = deque()
        x.append(self.root)
        while x:
            leaf = x.popleft()
            yield leaf.value
            if leaf.left:
                x.append(leaf.left)
            if leaf.right:
                x.append(leaf.right)

    def delete(self, val):
        self.root = self._delete(val, self.root)
        return None

    def _delete(self, val, leaf):
        def _descendants(leaf):
            if leaf.left:
                return _descendants(leaf.left)
            else:
                return leaf.value

        if not leaf:
            return None

        if leaf.value == val:
            self._size -= 1
            if leaf.left and leaf.right:
                leaf.value = _descendants(leaf.right)
                leaf.right = self._delete(leaf.value, leaf.right)
                return leaf
            elif leaf.left and not leaf.right:
                return leaf.left
            elif not leaf.left and leaf.right:
                return leaf.right
            else:
                return None

        elif leaf.value < val:
            if leaf.right:
                leaf.right = self._delete(val, leaf.right)
            return leaf

        else:
            if leaf.left:
                leaf.left = self._delete(val, leaf.left)
            return leaf

    def rotate_left(self, root):
        left = root.is_left()
        pivot = root.right_child

        if pivot is None:
            return
        root.right_child = pivot.left_child
        if pivot.left_child is not None:
            root.right_child.parent = root
        pivot.left_child = root
        pivot.parent = root.parent
        root.parent = pivot
        if left is None:
            self.root = pivot
        elif left:
            pivot.parent.left_child = pivot
        else:
            pivot.parent.right_child = pivot
        root.update_height(False)
        pivot.update_height(False)

    def rotate_right(self, root):
        left = root.is_left()
        pivot = root.left_child
        if pivot is None:
            return
        root.left_child = pivot.right_child
        if pivot.right_child is not None:
            root.left_child.parent = root
        pivot.right_child = root
        pivot.parent = root.parent
        root.parent = pivot
        if left is None:
            self.root = pivot
        elif left:
            pivot.parent.left_child = pivot
        else:
            pivot.parent.right_child = pivot
        root.update_height(False)
        pivot.update_height(False)

    def find_leftmost(self, node):
        if node.left_child is None:
            return node
        else:
            return self.find_leftmost(node.left_child)

    def find_rightmost(self, node):
        if node.right_child is None:
            return node
        else:
            return self.find_rightmost(node.right_child)

    def find_next(self, value):
        node = self.find(value)
        if (node is None) or (node.value != value):
            return None
        else:
            right_child = node.right_child
            if right_child is not None:
                node = self.find_leftmost(right_child)
            else:
                parent = node.parent
                while(parent is not None):
                    if node is parent.left_child:
                        break
                    node = parent
                    parent = node.parent
                node = parent
            if node is None:
                return node
            else:
                return node.value

    def find_prev(self, value):
        node = self.find(value)
        if (node is None) or (node.value != value):
            return None
        else:
            left_child = node.left_child
            if left_child is not None:
                node = self.find_leftmost(left_child)
            else:
                parent = node.parent
                while(parent is not None):
                    if node is parent.right_child:
                        break
                    node = parent
                    parent = node.parent
                node = parent
            if node is None:
                return node
            else:
                return node.value

    def balance(self, node):
        ''' There are four posabilities for rotation
            left-left=LL right-right=RR
            left-right=LR right-left=RL'''
        node.update_height(False)
        if node.balance == 2:
            if node.left_child.balance != -1:
                # LL rotation
                self.rotate_right(node)
                if node.parent.parent is not None:
                    self.balance(node.parent.parent)
            else:
                # LR rotation
                self.rotate_left(node.left_child)
                self.balance(node)
        elif node.balance == -2:
            if node.right_child.balance != 1:
                # RR rotation
                self.rotate_left(node)
                if node.parent.parent is not None:
                    self.balance(node.parent.parent)

            else:
                # RL rotation
                self.rotate_right(node.right_child)
                self.balance(node)
        else:
            if node.parent is not None:
                self.balance(node.parent)

    def sort(tree_maker, ascending=True):
        t = BST()
        for item in tree_maker:
            t.insert(item)
        ret = []
        if ascending:
            node = t.

if __name__ == '__main__':

    # x = range(100)
    # bst = BST()
    # for i in x:
    #     bst.insert(i)
    # bst.insert(42.1)
    # dot_graph = bst.get_dot()
    # t = subprocess.Popen(["dot", "-Tpng"], stdin=subprocess.PIPE)
    # t.communicate(dot_graph)

    def easy_tree():
        x = random.sample(range(100), 100)
        bst = BST()
        bst.insert(50)
        for i in x:
            bst.insert(i)
        bst.insert(42.1)
        bst.contains(42.1)

    def hard_tree():
        x = range(100)
        bst = BST()
        for i in x:
            bst.insert(i)
        bst.insert(42.1)
        bst.contains(42.1)

    print(timeit.Timer("easy_tree()", setup="from __main__ import easy_tree").timeit(number=1000))
    print(timeit.Timer("hard_tree()", setup="from __main__ import hard_tree").timeit(number=1000))
<|endoftext|>


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

def recv_from_sock(sock, nbytes):
    n = nbytes
    recieved = b''
    while n:
        bs = sock.recv(n)
        recieved += bs
        n -= len(bs)
    return recieved

def asc_dump(s):
    r = ''
    for c in s:
        r += chr(c) if (c >= 32 and c < 128) else '.'
    if r:
        print('[' + r + ']')

def _recv_from_server(server_sock):
    server_head = recv_from_sock(server_sock, 5)
    server_code = server_head[0]
    server_ln = int.from_bytes(server_head[1:], byteorder='big') -4
    server_data = recv_from_sock(server_sock, server_ln)
    print('S->C', chr(server_code), binascii.b2a_hex(server_data).decode('ascii'), end='')
    asc_dump(server_data)
    return server_head, server_data

def _recv_from_client(client_sock):
    client_head = recv_from_sock(client_sock, 5)
    client_code = client_head[0]
    client_ln = int.from_bytes(client_head[1:], byteorder='big')
    client_data = recv_from_sock(client_sock, client_ln-4)
    print('C->S', chr(client_code), binascii.b2a_hex(client_data).decode('ascii'), end='')
    asc_dump(client_data)
    return client_head, client_data

def _recv_from_client_tail(server_sock, client_sock):
    client_tail = recv_from_sock(client_sock, 5)
    assert client_tail == b'H\x00\x00\x00\x04'
    return client_tail

def parse_message(server_sock, client_sock):
    while True:
        server_head, server_data = _recv_from_server(server_sock)
        client_sock.send(server_head + server_data)
        if server_head[0] == 90:
            break

def read_login_packet(sock):
    head = recv_from_sock(sock, 4)
    ln = int.from_bytes(head, byteorder='big')
    return head + recv_from_sock(sock, ln-4)

def parse_login_response(server_sock, client_sock):
    server_head, server_data = _recv_from_server(server_sock)

    auth = int.from_bytes(server_data[:4], byteorder='big')
    if auth == 0:
        print('trust')
        client_sock.send(server_head + server_data)
    elif auth == 5:
        salt = server_data[4:]
        print('md5', binascii.b2a_hex(salt).decode('ascii'))
        client_sock.send(server_head + server_data)
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)

    elif auth == 10:
        print('SASL')
        client_sock.send(server_head + server_data)
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)

        server_head, server_data = _recv_from_server(server_sock)
        auth2 = int.from_bytes(server_data[:4], byteorder='big')
        assert auth2 == 11
        print('SCRAMFirst')
        client_sock.send(server_head + server_data)
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)

        server_head, server_data = _recv_from_server(server_sock)
        auth3 = int.from_bytes(server_data[:4], byteorder='big')
        assert auth3 == 12
        print('SCRAMFinal')
        client_sock.send(server_head + server_data)

        server_head, server_data = _recv_from_server(server_sock)
        # accept
        assert int.from_bytes(server_data[:4], byteorder='big') == 0
        client_sock.send(server_head + server_data)

def proxy_wire(server_name, server_port, listen_host, listen_port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind((listen_host, listen_port))
    sock.listen(1)
    client_sock, addr = sock.accept()
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sock.connect((server_name, server_port))

    login_packet = read_login_packet(client_sock)
    print('C->S login ', binascii.b2a_hex(login_packet).decode('ascii'), end='')
    asc_dump(login_packet)
    server_sock.send(login_packet)
    parse_login_response(server_sock, client_sock)

    parse_message(server_sock, client_sock)

    while True:
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)
        if chr(client_head[0]) == 'X':
            break
        _recv_from_client_tail(server_sock, client_sock)

        parse_message(server_sock, client_sock)

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage : ' + sys.argv[0] + ' server[:port] [listen_host:]listen_port')
        sys.exit()

    server = sys.argv[1].split(':')
    server_name = server[0]
    if len(server) == 1:
        server_port = 5432
    else:
        server_port = int(server[1])

    listen = sys.argv[2].split(':')
    if len(listen) == 1:
        listen_host = 'localhost'
        listen_port = int(listen[0])
    else:
        listen_host = listen[0]
        listen_port = int(listen[1])

    proxy_wire(server_name, server_port, listen_host, listen_port)
<commit_msg>parse accept data (proxy tool)<commit_after>#!/usr/bin/env python3
##############################################################################
#The MIT License (MIT)
#
#Copyright (c) 2016 Hajime Nakagami
#
#Permission is hereby granted, free of charge, to any person obtaining a copy
#of this software and associated documentation files (the "Software"), to deal
#in the Software without restriction, including without limitation the rights
#to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#copies of the Software, and to permit persons to whom the Software is
#furnished to do so, subject to the following conditions:
#
#The above copyright notice and this permission notice shall be included in all
#copies or substantial portions of the Software.
#
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#SOFTWARE.
##############################################################################
import sys
import socket
import binascii

def recv_from_sock(sock, nbytes):
    n = nbytes
    recieved = b''
    while n:
        bs = sock.recv(n)
        recieved += bs
        n -= len(bs)
    return recieved

def asc_dump(s):
    r = ''
    for c in s:
        r += chr(c) if (c >= 32 and c < 128) else '.'
    if r:
        print('[' + r + ']')

def _recv_from_server(server_sock):
    server_head = recv_from_sock(server_sock, 5)
    server_code = server_head[0]
    server_ln = int.from_bytes(server_head[1:], byteorder='big') -4
    server_data = recv_from_sock(server_sock, server_ln)
    print('S->C', chr(server_code), binascii.b2a_hex(server_data).decode('ascii'), end='')
    asc_dump(server_data)
    return server_head, server_data

def _recv_from_client(client_sock):
    client_head = recv_from_sock(client_sock, 5)
    client_code = client_head[0]
    client_ln = int.from_bytes(client_head[1:], byteorder='big')
    client_data = recv_from_sock(client_sock, client_ln-4)
    print('C->S', chr(client_code), binascii.b2a_hex(client_data).decode('ascii'), end='')
    asc_dump(client_data)
    return client_head, client_data

def _recv_from_client_tail(server_sock, client_sock):
    client_tail = recv_from_sock(client_sock, 5)
    assert client_tail == b'H\x00\x00\x00\x04'
    return client_tail

def parse_message(server_sock, client_sock):
    while True:
        server_head, server_data = _recv_from_server(server_sock)
        client_sock.send(server_head + server_data)
        if server_head[0] == 90:
            break

def read_login_packet(sock):
    head = recv_from_sock(sock, 4)
    ln = int.from_bytes(head, byteorder='big')
    return head + recv_from_sock(sock, ln-4)

def parse_login_response(server_sock, client_sock):
    server_head, server_data = _recv_from_server(server_sock)

    auth = int.from_bytes(server_data[:4], byteorder='big')
    if auth == 0:
        print('trust')
        client_sock.send(server_head + server_data)
    elif auth == 5:
        salt = server_data[4:]
        print('md5', binascii.b2a_hex(salt).decode('ascii'))
        client_sock.send(server_head + server_data)
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)

        server_head, server_data = _recv_from_server(server_sock)
        # accept
        assert int.from_bytes(server_data[:4], byteorder='big') == 0
        client_sock.send(server_head + server_data)
    elif auth == 10:
        print('SASL')
        client_sock.send(server_head + server_data)
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)

        server_head, server_data = _recv_from_server(server_sock)
        auth2 = int.from_bytes(server_data[:4], byteorder='big')
        assert auth2 == 11
        print('SCRAMFirst')
        client_sock.send(server_head + server_data)
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)

        server_head, server_data = _recv_from_server(server_sock)
        auth3 = int.from_bytes(server_data[:4], byteorder='big')
        assert auth3 == 12
        print('SCRAMFinal')
        client_sock.send(server_head + server_data)

        server_head, server_data = _recv_from_server(server_sock)
        # accept
        assert int.from_bytes(server_data[:4], byteorder='big') == 0
        client_sock.send(server_head + server_data)

def proxy_wire(server_name, server_port, listen_host, listen_port):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind((listen_host, listen_port))
    sock.listen(1)
    client_sock, addr = sock.accept()
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sock.connect((server_name, server_port))

    login_packet = read_login_packet(client_sock)
    print('C->S login ', binascii.b2a_hex(login_packet).decode('ascii'), end='')
    asc_dump(login_packet)
    server_sock.send(login_packet)
    parse_login_response(server_sock, client_sock)

    parse_message(server_sock, client_sock)

    while True:
        client_head, client_data = _recv_from_client(client_sock)
        server_sock.send(client_head + client_data)
        if chr(client_head[0]) == 'X':
            break
        _recv_from_client_tail(server_sock, client_sock)

        parse_message(server_sock, client_sock)

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage : ' + sys.argv[0] + ' server[:port] [listen_host:]listen_port')
        sys.exit()

    server = sys.argv[1].split(':')
    server_name = server[0]
    if len(server) == 1:
        server_port = 5432
    else:
        server_port = int(server[1])

    listen = sys.argv[2].split(':')
    if len(listen) == 1:
        listen_host = 'localhost'
        listen_port = int(listen[0])
    else:
        listen_host = listen[0]
        listen_port = int(listen[1])

    proxy_wire(server_name, server_port, listen_host, listen_port)
<|endoftext|>