Python config.OUTPUT_DIR Examples

def process_subdomain(self):
        helper.install_domains()
        sqlitepath = os.path.join(config.OUTPUT_DIR,'domains.db')
        conn = db.connect(sqlitepath)
        conn.text_factory = str
        cursor = conn.cursor()
        sql = "INSERT INTO domains(domain, ip, cname, cdn, internal) VALUES(?, ?, ?, ?, ?)"

        ips = set()
        cdn_ip = set()

        self.subdomains.update(self.domains)

        for domain in self.subdomains:
            cname = tools.get_cname(domain)
            cdn = tools.get_cdn(domain, cname)
            ipl = self.domain_ip.get(domain, None)
            if cdn:
                self.cdn_domain.add(domain)
            if not ipl:
                ipl = tools.resolve_host_ip(domain)
            else:
                ipl = ipl.split(",")
            for ip in ipl:
                internal = tools.is_internal_ip(ip)
                if not cdn and not internal:
                    ips.add(ip)
                elif cdn:
                    self.takeover_domain_check.add((domain, ip, cname))
                    cdn_ip.add(ip)
                if not internal:
                    self.internal_domain.add(domain)
                try:
                    status = cursor.execute(sql, (domain, ip, cname, cdn, internal))
                    conn.commit()
                except Exception as e:
                    print e
        self.ips = ips-cdn_ip
        with open(os.path.join(config.OUTPUT_DIR,config.IPS), 'w') as f:
            f.write('\n'.join(self.ips).strip()) 

def get_scores(self, preds, save=False):
        preds = [label_path(self.id2type[x]) for x in preds]
        def vec2type(v):
            s = []
            for i in range(len(v)):
                if v[i]:
                    s.extend(label_path(self.id2type[i]))
            return set(s)
        labels_test = [vec2type(x) for x in self.labels_test]
        if save:
            labels_test = [vec2type(x) for x in self.labels]
        acc = strict(labels_test, preds)
        _, _, macro = loose_macro(labels_test, preds)
        _, _, micro = loose_micro(labels_test, preds)

        if save:
            outfile = open(os.path.join(config.OUTPUT_DIR, self.__str__() + ".tsv"), "w")
            for x, y in zip(preds, labels_test):
                t1 = "|".join(list(x))
                t2 = "|".join(list(y))
                outfile.write(t1 + "\t" + t2 + "\n")
            outfile.close()

        return acc, macro, micro 

def install_domains():
	sqlitepath = os.path.join(config.OUTPUT_DIR, "domains.db")
	install = ''
	if not os.path.exists(sqlitepath):
		install = '''
	CREATE TABLE `domains`(
		`domain` varchar(255) NOT NULL,
		`ip` TEXT NOT NULL,
		`cname` varchar(255),
		`cdn` INTEGER,
		`internal` INTEGER,
		PRIMARY KEY(`domain`, `ip`)
	);
	'''
	if install:
		conn = get_domains_conn()
		cursor = conn.cursor()
		cursor.execute(install)
		conn.commit()
		conn.close() 

def install_ports():
	sqlitepath = os.path.join(config.OUTPUT_DIR, "ports.db")
	install = ''
	if not os.path.exists(sqlitepath):
		install = '''
	CREATE TABLE open(
		`ip` VARCHAR(64) NOT NULL,
		`port` INTEGER,
		`service` varchar(64),
		`comment` TEXT,
		PRIMARY KEY(`ip`, `port`)
	);
	'''
	if install:
		conn = conn = get_ports_conn()
		cursor = conn.cursor()
		cursor.execute(install)
		conn.commit()
		conn.close() 

def get_ports_conn():
	sqlitepath = os.path.join(config.OUTPUT_DIR, "ports.db")
	conn = db.connect(sqlitepath)
	conn.text_factory = str
	return conn 

def get_domains_conn():
	sqlitepath = os.path.join(config.OUTPUT_DIR, "domains.db")
	conn = db.connect(sqlitepath)
	conn.text_factory = str
	return conn 

def parse_domains_brute(domain, extip=None):
	'''
	如果域名泛解析，则通过HTTP请求的Host来判断是否真的绑定在webserver上
	在检查响应的时候，一般同一个错误页面的响应长度是一样的，除非响应中包含 host，所以需要在替换掉host之后再比较长度
	'''
	def get_error_page(extip, fhost):
		error_page = ''
		try:
			error_page = requests.get('https://%s' % extip, headers={'host': fhost}, verify=True).text.replace(fhost, "")
		except Exception as e:
			pass
		if not error_page:
			try:
				fhost = 'salt66666666.'+domain
				error_page = requests.get('http://%s' % extip, headers={'host': fhost}).text.replace(fhost, "")
			except Exception as e:
				pass
		return len(error_page)

	with open(os.path.join(config.OUTPUT_DIR, '%s.txt'%domain), 'r') as f:
		data = f.read().strip()
	ret = {}

	if extip:
		fhost = 'salt66666666.'+domain
		error_page = get_error_page(extip, fhost)

	for line in data.split('\n'):
		if not line.strip():
			continue
		line = line.replace(' ', '').replace('\t', '')
		parts = line.split(domain)
		if extip and extip in line:
			if not error_page:
				continue
			else:
				page = get_error_page(extip, parts[0]+domain)
				if page == error_page:
					continue

		ret[parts[0]+domain] = parts[1]
	return ret 

def report(self):
        json.dump(self.ip_all, open(os.path.join(config.OUTPUT_DIR, "ip_all.json"), "w"))
        json.dump(list(self.cdn_domain), open(os.path.join(config.OUTPUT_DIR, "cdn_domain.json"), "w"))
        json.dump(list(self.internal_domain), open(os.path.join(config.OUTPUT_DIR, "internal_domain.json"), "w"))
        json.dump(list(self.extensive_domain), open(os.path.join(config.OUTPUT_DIR, "extensive_domain.json"), "w"))

        with open(os.path.join(config.OUTPUT_DIR, 'domain_takeover.txt'), 'a') as f:
            f.write('\n'.join(self.takeover_domain).strip())
        tools.report(self.ip_all, outname=config.REPORT_FILENAME) 

def report_subdomain(self):
        domains = set()
        conn = helper.get_domains_conn()
        cur = conn.cursor()
        cur.execute("SELECT * FROM domains WHERE cdn=0 and internal=0")
        rows = cur.fetchall()
        for row in rows:
            domain, ip, cname, cdn, internal = row
            domains.add(domain)
        json.dump(list(domains), open(os.path.join(config.OUTPUT_DIR, "all_subdomains.json"), "w")) 

def check_valid(model):
    file = "%s/All/test.pred.%s.csv" % (config.OUTPUT_DIR, model)
    try:
        df = pd.read_csv(file)
        if df.shape[0] == config.TEST_SIZE:
            return True
        else:
            return False
    except:
        return False 

def active_search(self):
        scanable_domain = set()
        for d in self.subdomains:
            scanable_domain.update(tools.scanable_subdomain(d))

        self.subdomains = set(filter(lambda x: not x.startswith('*.'), self.subdomains))
        scanable_domain.update(self.domains)

        for domain in scanable_domain:
            isext, ip = tools.check_extensive_domain(domain)
            if isext:
                self.extensive_domain.add(domain)
            if not os.path.exists(os.path.join(config.OUTPUT_DIR, '%s.txt'%domain)):
                if tools.get_domain(domain) == domain:
                    d = subDomainsBrute.SubNameBrute(target=domain, options=subDomainsBruteOpt(domain))
                else:
                    d = subDomainsBrute.SubNameBrute(target=domain, options=subDomainsBruteOpt(domain, "next_sub.txt"))
                d.run()
                d.outfile.flush()
                d.outfile.close()
            r = helper.parse_domains_brute(domain, ip)
            self.subdomains.update(r.keys())
            self.domain_ip.update(r) 

def __init__(self, domain, dictionary="subnames.txt"):
        self.file= "subDomainsBrute"+os.sep+"dict"+os.sep+dictionary
        self.threads = 200
        self.output = os.path.join(config.OUTPUT_DIR, '%s.txt'%domain)
        self.i = False
        self.full_scan = False 

def create_directories():
    """Create all the directories in the /data directories which are used for preprocessing/training/evaluating."""

    directories = [TILES_DIR, WATER_BITMAPS_DIR, WGS84_DIR, LABELS_DIR, MODELS_DIR, OUTPUT_DIR, TENSORBOARD_DIR]
    for directory in directories:
        save_makedirs(directory) 

def refit(self):
        for i in range(self.n_iter):
            if self.refit_once and i >= 1:
                break
            X_train, y_train, X_test = self.feature._get_train_test_data(i)
            self.learner.fit(X_train, y_train)
            if i == 0:
                y_pred = self.learner.predict(X_test)
                if hasattr(self.learner.learner, "predict_proba"):
                    y_proba = self.learner.learner.predict_proba(X_test)
            else:
                y_pred += self.learner.predict(X_test)
                if hasattr(self.learner.learner, "predict_proba"):
                    y_proba += self.learner.learner.predict_proba(X_test)
        if not self.refit_once:
            y_pred /= float(self.n_iter)
            if hasattr(self.learner.learner, "predict_proba"):
                y_proba /= float(self.n_iter)

        id_test = self.feature.data_dict["id_test"].astype(int)

        # save
        fname = "%s/%s/test.pred.%s.csv"%(config.OUTPUT_DIR, "All", self.__str__())
        pd.DataFrame({"id": id_test, "prediction": y_pred}).to_csv(fname, index=False)
        if hasattr(self.learner.learner, "predict_proba"):
            fname = "%s/%s/test.proba.%s.csv"%(config.OUTPUT_DIR, "All", self.__str__())
            columns = ["proba%d"%i for i in range(y_proba.shape[1])]
            print(y_proba.shape)
            print(len(columns))
            df = pd.DataFrame(y_proba, columns=columns)
            df["id"] = id_test
            df.to_csv(fname, index=False)

        # submission
        fname = "%s/test.pred.%s.[Mean%.6f]_[Std%.6f].csv"%(
            config.SUBM_DIR, self.__str__(), self.rmse_cv_mean, self.rmse_cv_std)
        pd.DataFrame({"id": id_test, "relevance": y_pred}).to_csv(fname, index=False)
        return self 

def refit(self):
        X_train, y_train, X_test = self.feature._get_train_test_data()
        if self.plot_importance:
            feature_names = self.feature._get_feature_names()
            self.learner.fit(X_train, y_train, feature_names)
            y_pred = self.learner.predict(X_test, feature_names)
        else:
            self.learner.fit(X_train, y_train)
            y_pred = self.learner.predict(X_test)
        
        id_test = self.feature.data_dict["id_test"].astype(int)

        # save
        fname = "%s/%s/test.pred.%s.csv"%(config.OUTPUT_DIR, "All", self.__str__())
        pd.DataFrame({"id": id_test, "prediction": y_pred}).to_csv(fname, index=False)
        if hasattr(self.learner.learner, "predict_proba"):
            if self.plot_importance:
                feature_names = self.feature._get_feature_names()
                y_proba = self.learner.learner.predict_proba(X_test, feature_names)
            else:
                y_proba = self.learner.learner.predict_proba(X_test)
            fname = "%s/%s/test.proba.%s.csv"%(config.OUTPUT_DIR, "All", self.__str__())
            columns = ["proba%d"%i for i in range(y_proba.shape[1])]
            df = pd.DataFrame(y_proba, columns=columns)
            df["id"] = id_test
            df.to_csv(fname, index=False)

        # submission
        fname = "%s/test.pred.%s.[Mean%.6f]_[Std%.6f].csv"%(
            config.SUBM_DIR, self.__str__(), self.rmse_cv_mean, self.rmse_cv_std)
        pd.DataFrame({"id": id_test, "relevance": y_pred}).to_csv(fname, index=False)

        # plot importance
        if self.plot_importance:
            ax = self.learner.plot_importance()
            ax.figure.savefig("%s/%s.pdf"%(config.FIG_DIR, self.__str__()))
        return self 

def cv(self):
        start = time.time()
        if self.verbose:
            self.logger.info("="*50)
            self.logger.info("Task")
            self.logger.info("      %s" % str(self.__str__()))
            self.logger.info("Param")
            self._print_param_dict(self.learner.param_dict)
            self.logger.info("Result")
            self.logger.info("      Run      RMSE        Shape")
    
        rmse_cv = np.zeros(self.n_iter)
        for i in range(self.n_iter):
            # data
            X_train, y_train, X_valid, y_valid = self.feature._get_train_valid_data(i)
            # fit
            self.learner.fit(X_train, y_train)
            y_pred = self.learner.predict(X_valid)
            rmse_cv[i] = dist_utils._rmse(y_valid, y_pred)
            # log
            self.logger.info("      {:>3}    {:>8}    {} x {}".format(
                i+1, np.round(rmse_cv[i],6), X_train.shape[0], X_train.shape[1]))
            # save
            fname = "%s/Run%d/valid.pred.%s.csv"%(config.OUTPUT_DIR, i+1, self.__str__())
            df = pd.DataFrame({"target": y_valid, "prediction": y_pred})
            df.to_csv(fname, index=False, columns=["target", "prediction"])
            if hasattr(self.learner.learner, "predict_proba"):
                y_proba = self.learner.learner.predict_proba(X_valid)
                fname = "%s/Run%d/valid.proba.%s.csv"%(config.OUTPUT_DIR, i+1, self.__str__())
                columns = ["proba%d"%i for i in range(y_proba.shape[1])]
                df = pd.DataFrame(y_proba, columns=columns)
                df["target"] = y_valid
                df.to_csv(fname, index=False)

        self.rmse_cv_mean = np.mean(rmse_cv)
        self.rmse_cv_std = np.std(rmse_cv)
        end = time.time()
        _sec = end - start
        _min = int(_sec/60.)
        if self.verbose:
            self.logger.info("RMSE")
            self.logger.info("      Mean: %.6f"%self.rmse_cv_mean)
            self.logger.info("      Std: %.6f"%self.rmse_cv_std)
            self.logger.info("Time")
            if _min > 0:
                self.logger.info("      %d mins"%_min)
            else:
                self.logger.info("      %d secs"%_sec)
            self.logger.info("-"*50)
        return self 

def main(options):

    # create sub folder
    subm_folder = "%s/ensemble_selection"%config.SUBM_DIR
    os_utils._create_dirs( [subm_folder] )
    subm_prefix = "%s/test.pred.[%s]" % (subm_folder, options.outfile)

    # get model list
    log_folder = "%s/level%d_models"%(config.LOG_DIR, options.level-1)
    model_list = get_model_list(log_folder, options.size)

    # get instance splitter
    if options.level not in [2, 3]:
        inst_splitter = None
    elif options.level == 2:
        inst_splitter = splitter_level2
    elif options.level == 3:
        inst_splitter = splitter_level3

    ees = ExtremeEnsembleSelection(
            model_folder=config.OUTPUT_DIR, 
            model_list=model_list, 
            subm_prefix=subm_prefix, 
            weight_opt_max_evals=options.weight_opt_max_evals, 
            w_min=-1., 
            w_max=1., 
            inst_subsample=options.inst_subsample,
            inst_subsample_replacement=options.inst_subsample_replacement,
            inst_splitter=inst_splitter,
            model_subsample=options.model_subsample,
            model_subsample_replacement=options.model_subsample_replacement,
            bagging_size=options.bagging_size, 
            init_top_k=options.init_top_k,
            epsilon=options.epsilon,
            multiprocessing=False, 
            multiprocessing_num_cores=config.NUM_CORES,
            enable_extreme=options.enable_extreme,
            random_seed=config.RANDOM_SEED
        )
    ees.go()