示例#1
0
func TestMain(m *testing.M) {
	err := Init()
	if err != nil {
		panic(err)
	}
	clients := db.Clients{
		DB:       sqlDB,
		Internet: db.NewInternetClient(redisPool),
		Maxmind:  db.NewMaxmindClient(mmCountryDB, mmCityDB),
	}
	mux, err := apiMux(clients)
	if err != nil {
		panic(err)
	}

	ts := httptest.NewServer(mux)
	testclient = client.NewClient(ts.URL, nil)
	defer ts.Close()
	os.Exit(m.Run())
}
示例#2
0
// Run runs the initialized server.
func Run() {
	var wg sync.WaitGroup

	// start monitor server
	go startMonitoring(*monitorBindAddr)

	// start the getpublic ip updater.
	go func() {
		_ = shared.GetPublicIPAddr()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		lg.V(2).Infoln("Loading recent sessions from postgres...")
		recents, err := sqlDB.RecentSuggestionSessions(20000)
		if err != nil {
			lg.Fatal(err)
		}
		db.SessionTokens.Reset(recents)
		lg.V(2).Infof("Loaded %d sessions from postgres...", len(recents))
		lg.Flush()

	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		conn := redisPool.Get()
		defer conn.Close()
		lg.V(2).Infoln("BGPDump refresh started...")
		n, err := internet.RefreshBGPDump(conn)
		lg.V(2).Infof("BGPDump refresh ended, %d items added.", n)
		lg.Flush()
		if err != nil {
			if *offline {
				lg.Infoln("offline", err)
			} else {
				lg.Fatal(err)
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		conn := redisPool.Get()
		defer conn.Close()
		lg.V(2).Infoln("CIDRReport refresh started...")
		n, err := internet.RefreshCIDRReport(conn)
		lg.V(2).Infof("CIDRReport refresh ended, %d items added", n)
		if err != nil {
			if *offline {
				lg.Infoln("offline", err)
			} else {
				lg.Fatal(err)
			}
		}
	}()
	wg.Wait()

	// start signal handling
	wg.Add(1)
	go func() {
		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGINT)
		lg.Infoln(<-ch)
		wg.Done()
	}()

	internetClient := db.NewInternetClient(redisPool)
	maxmindClient := db.NewMaxmindClient(mmCountryDB, mmCityDB)

	clients := db.Clients{
		DB:       sqlDB,
		Internet: internetClient,
		Maxmind:  maxmindClient,
	}

	// start http json api server
	go func(addr string, dba db.Clients) {
		mux, err := apiMux(dba)
		lg.Info("Starting http server", addr)
		err = http.ListenAndServe(addr, mux)
		if err != nil {
			lg.Fatal(err)
		}
	}(*apiBindAddr, clients)

	// start http export api server
	go func(addr string, dba db.Clients) {
		if *exportApiSecretKey == "" {
			lg.Warningln("exportApiSecretKey flag/env not set, will not start export api server")
			b := make([]byte, 32)
			_, err := rand.Read(b)
			if err != nil {
				lg.Fatalf("random generator not functioning...")
				return
			}
			suggestedkey := base64.StdEncoding.EncodeToString(b)
			lg.Infoln("suggested export key:", suggestedkey)
			return
		}
		key, err := base64.StdEncoding.DecodeString(*exportApiSecretKey)
		if err != nil {
			lg.Fatalf("could not decode export api secret key: %s", *exportApiSecretKey)
		}

		mux, err := apiMuxExport(dba, key)
		lg.Info("Starting export api server", addr)
		err = http.ListenAndServe(addr, mux)
		if err != nil {
			lg.Fatal(err)
		}
	}(*exportApiBindAddr, clients)

	go analysis.StartAnalysis(clients)
	startMeasurer(clients)

	wg.Wait()
}