示例#1
0
文件: main.go 项目: jcjones/ct-sql
func main() {
	log.SetFlags(0)
	log.SetPrefix("")
	dbConnectStr, err := sqldb.RecombineURLForDB(*config.DbConnect)
	if err != nil {
		log.Printf("unable to parse %s: %s", *config.DbConnect, err)
	}

	if len(dbConnectStr) == 0 || (config.CensysPath == nil && config.LogUrl == nil) {
		config.Usage()
		os.Exit(2)
	}

	db, err := sql.Open("mysql", dbConnectStr)
	if err != nil {
		log.Fatalf("unable to open SQL: %s: %s", dbConnectStr, err)
	}
	if err = db.Ping(); err != nil {
		log.Fatalf("unable to ping SQL: %s: %s", dbConnectStr, err)
	}

	var certFolderDB *utils.FolderDatabase
	if config.CertPath != nil && len(*config.CertPath) > 0 {
		certFolderDB, err = utils.NewFolderDatabase(*config.CertPath, 0444, *config.CertsPerFolder)
		if err != nil {
			log.Fatalf("unable to open Certificate Path: %s: %s", config.CertPath, err)
		}
	}

	dialect := gorp.MySQLDialect{Engine: "InnoDB", Encoding: "UTF8"}
	dbMap := &gorp.DbMap{Db: db, Dialect: dialect}
	entriesDb := &sqldb.EntriesDatabase{DbMap: dbMap,
		Verbose: *config.Verbose, FullCerts: certFolderDB,
		KnownIssuers: make(map[string]int)}
	err = entriesDb.InitTables()
	if err != nil {
		log.Fatalf("unable to prepare SQL: %s: %s", dbConnectStr, err)
	}

	if config.LogUrl != nil && len(*config.LogUrl) > 5 {
		ctLogUrl, err := url.Parse(*config.LogUrl)
		if err != nil {
			log.Fatalf("unable to set Certificate Log: %s", err)
		}

		ctLog := client.New(*config.LogUrl, nil)

		log.Printf("Starting download from log %s, fullCerts=%t\n", ctLogUrl, (certFolderDB != nil))

		err = downloadLog(ctLogUrl, ctLog, entriesDb)
		if err != nil {
			log.Fatalf("error while updating CT entries: %s", err)
		}

		os.Exit(0)
	}

	var importer censysdata.Importer
	if config.CensysUrl != nil && len(*config.CensysUrl) > 5 {
		urlImporter, err := censysdata.OpenURL(*config.CensysUrl)
		if err != nil {
			log.Fatalf("unable to open Censys URL: %s", err)
		}
		importer = urlImporter
	} else if config.CensysPath != nil && len(*config.CensysPath) > 5 {
		fileImporter, err := censysdata.OpenFile(*config.CensysPath)
		if err != nil {
			log.Fatalf("unable to open Censys file: %s", err)
		}
		importer = fileImporter
		defer fileImporter.Close()
	} else if *config.CensysStdin {
		stdinImporter, err := censysdata.OpenFileHandle(os.Stdin)
		if err != nil {
			log.Fatalf("unable to open stdin: %s", err)
		}
		importer = stdinImporter
		defer stdinImporter.Close()
	}

	if importer != nil {
		log.Printf("Starting Censys Import, using %s, fullCerts=%t\n", importer.String(), (certFolderDB != nil))

		wg := new(sync.WaitGroup)
		err = processImporter(importer, entriesDb, wg)

		if err != nil {
			log.Fatalf("error while running importer: %s", err)
		}

		wg.Wait()
		os.Exit(0)
	}

	// Didn't include a mandatory action, so print usage and exit.
	config.Usage()
	os.Exit(2)
}
示例#2
0
文件: main.go 项目: jcjones/ct-sql
func processImporter(importer censysdata.Importer, db *sqldb.EntriesDatabase, wg *sync.WaitGroup) error {
	entryChan := make(chan censysdata.CensysEntry, 100)
	defer close(entryChan)
	statusChan := make(chan OperationStatus, 1)
	defer close(statusChan)
	wg.Add(1)
	defer wg.Done()

	displayProgress(statusChan, wg)
	numWorkers := *config.NumThreads * runtime.NumCPU()
	for i := 0; i < numWorkers; i++ {
		go insertCensysWorker(entryChan, db, wg)
	}

	startOffset := *config.OffsetByte

	//
	// Fast forward
	//
	if *config.Offset > 0 && *config.OffsetByte > 0 {
		return fmt.Errorf("You may not set both offset and offsetByte")
	}

	if *config.Offset > 0 {
		err := importer.SeekLine(*config.Offset)
		if err != nil {
			return err
		}
	}

	if *config.OffsetByte > 0 {
		err := importer.SeekByte(*config.OffsetByte)
		if err != nil {
			return err
		}
	}

	var maxOffset uint64

	log.Printf("Starting import from offset=%d, line_limit=%d", importer.ByteOffset(), *config.Limit)

	// We've already fast-forwarded, so start at 0.
	for count := uint64(0); ; count++ {
		if *config.Limit > uint64(0) && count >= *config.Limit {
			return nil
		}
		ent, err := importer.NextEntry()
		if err != nil || ent == nil {
			return err
		}
		if count%128 == 0 {
			// Lazily obtain the max offset
			if maxOffset < 1 {
				maxOffset, err = importer.Size()
				if err != nil {
					return err
				}
			}

			statusChan <- OperationStatus{int64(startOffset), int64(ent.Offset), int64(maxOffset)}
		}
		entryChan <- *ent
	}

	return nil
}