Prepare for API key

Client.go

@@ -2,13 +2,15 @@ package LeakIXClient
 
 import (
 	"encoding/json"
+	"errors"
 	"fmt"
 	"github.com/gorilla/websocket"
 	"gitlab.nobody.run/tbi/core"
 	"io/ioutil"
 	"log"
 	"net/http"
 	url2 "net/url"
+	"strings"
 	"time"
 )
 
@@ -29,6 +31,17 @@ type SearchResultsClient struct {
 	SearchResults []SearchResult
 	Position      int
 	Page          int
+	ApiKey        string
+	Endpoint      string
+}
+
+const defaultEndpoint = "https://leakix.net"
+
+func (sc *SearchResultsClient) GetEndpoint() string {
+	if len(sc.Endpoint) > 8 {
+		return sc.Endpoint
+	}
+	return defaultEndpoint
 }
 
 func (sc *SearchResultsClient) Next() bool {
@@ -37,7 +50,7 @@ func (sc *SearchResultsClient) Next() bool {
 		return true
 	}
 	// Try to load next page
-	results, _ := GetSearchResults(sc.Scope, sc.Query, sc.Page)
+	results, _ := sc.GetSearchResults(sc.Scope, sc.Query, sc.Page)
 	for _, result := range results {
 		sc.SearchResults = append(sc.SearchResults, result)
 	}
@@ -53,12 +66,13 @@ func (sc *SearchResultsClient) SearchResult() SearchResult {
 	return sc.SearchResults[sc.Position-1]
 }
 
-func GetSearchResults(scope string, query string, page int) ([]SearchResult, error) {
+func (sc *SearchResultsClient) GetSearchResults(scope string, query string, page int) ([]SearchResult, error) {
 	url := fmt.Sprintf(
-		"https://leakix.net/search?scope=%s&q=%s&page=%d", url2.QueryEscape(scope), url2.QueryEscape(query), page)
+		"%s/search?scope=%s&q=%s&page=%d", sc.GetEndpoint(), url2.QueryEscape(scope), url2.QueryEscape(query), page)
 	var searchResults []SearchResult
 	req, _ := http.NewRequest("GET", url, nil)
 	req.Header.Set("Accept", "application/json")
+	req.Header.Set("api-key", sc.ApiKey)
 	resp, err := HttpClient.Do(req)
 	if err != nil {
 		return searchResults, err
@@ -75,15 +89,24 @@ func GetSearchResults(scope string, query string, page int) ([]SearchResult, err
 	return searchResults, nil
 }
 
-func GetChannel(scope string) (chan SearchResult, error) {
+func (sc *SearchResultsClient) GetChannel(scope string) (chan SearchResult, error) {
 	channel := make(chan SearchResult)
-	wsConnection, _, err := websocket.DefaultDialer.Dial("wss://leakix.net/ws/" + scope, map[string][]string{"Origin":{"leakix.net:443"}})
+	endpointUrl, err := url2.Parse(sc.GetEndpoint())
+	if err != nil {
+		return nil, errors.New("invalid endpoint")
+	}
+	endpointUrl.Scheme = strings.Replace(endpointUrl.Scheme, "http", "ws", -1)
+	log.Println(endpointUrl.String())
+	wsConnection, _, err := websocket.DefaultDialer.Dial(endpointUrl.String()+"/ws/"+scope, map[string][]string{
+		"Origin": {endpointUrl.Host + ":" + endpointUrl.Port()},
+		"api-key":{sc.ApiKey},
+	})
 	if err != nil {
 		return nil, err
 	}
 	go func() {
 		searchResult := SearchResult{}
-		for  {
+		for {
 			err := wsConnection.ReadJSON(&searchResult)
 			if err != nil {
 				log.Println("Error parsing websocket results. Is your scope correct?")
@@ -93,4 +116,4 @@ func GetChannel(scope string) (chan SearchResult, error) {
 		}
 	}()
 	return channel, nil
-}
+}

cmd/leakix-ns/main.go

@@ -15,6 +15,7 @@ func main() {
 	flag.StringVar(&app.Domain, "d", "", "Specify domain")
 	flag.BoolVar(&app.OutputJson, "j", false, "JSON mode, (excludes -t)")
 	flag.IntVar(&app.Limit, "l", 100, "Limit results output")
+	flag.StringVar(&app.ApiKey, "k", "", "API Key")
 	flag.Usage = func() {
 		fmt.Printf("Usage of leakix-dns: \n")
 		fmt.Printf("  ./leakix -d <domain> -l 200\n\n")
@@ -35,12 +36,15 @@ type App struct {
 	Searcher   *LeakIXClient.SearchResultsClient
 	Reverse    map[string][]LeakIXClient.SearchResult
 	Forward    map[string][]LeakIXClient.SearchResult
+	ApiKey     string
 }
 
 func (app *App) Run() {
 	app.Searcher = &LeakIXClient.SearchResultsClient{
-		Scope: "service",
-		Query: fmt.Sprintf("hostname:\"%s\" OR reverse:\"%s\" OR ip:\"%s\"", app.Domain, app.Domain, app.Domain),
+		Scope:    "service",
+		Query:    fmt.Sprintf("hostname:\"%s\" OR reverse:\"%s\" OR ip:\"%s\"", app.Domain, app.Domain, app.Domain),
+		ApiKey:   app.ApiKey,
+		Endpoint: "https://staging.leakix.net",
 	}
 	app.Reverse = make(map[string][]LeakIXClient.SearchResult)
 	app.Forward = make(map[string][]LeakIXClient.SearchResult)

cmd/leakix/main.go

@@ -16,6 +16,8 @@ var scope string
 var query string
 var liveStream bool
 var outputTemplate string
+var apiKey string
+var endPoint string
 var limit int
 var tmpl *template.Template
 var err error
@@ -28,6 +30,9 @@ func main() {
 	flag.BoolVar(&liveStream, "r", false, "Realtime mode, (excludes -q)")
 	flag.BoolVar(&outputJson, "j", false, "JSON mode, (excludes -t)")
 	flag.StringVar(&outputTemplate, "t", "{{ .Ip }}:{{ .Port }}", "Specify output template")
+	flag.StringVar(&apiKey, "k", "", "Specify API key")
+	flag.StringVar(&endPoint, "e", "https://leakix.net", "Leakix endpoint to use")
+
 	flag.IntVar(&limit, "l", 100, "Limit results output")
 	flag.Usage = func() {
 		fmt.Printf("Usage of leakix: \n")
@@ -53,8 +58,10 @@ func main() {
 
 func Search() {
 	searcher := LeakIXClient.SearchResultsClient{
-		Scope:         scope,
-		Query:         query,
+		Scope:    scope,
+		Query:    query,
+		ApiKey:   apiKey,
+		Endpoint: endPoint,
 	}
 	count := 0
 	for searcher.Next() {
@@ -68,13 +75,17 @@ func Search() {
 
 func LiveStream() {
 	count := 0
-	serviceChannel, err := LeakIXClient.GetChannel(scope)
+	searcher := LeakIXClient.SearchResultsClient{
+		ApiKey:   apiKey,
+		Endpoint: endPoint,
+	}
+	serviceChannel, err := searcher.GetChannel(scope)
 	if err != nil {
 		log.Println("Websocket connection error:")
 		log.Fatal(err)
 	}
 	for {
-		service := <- serviceChannel
+		service := <-serviceChannel
 		count++
 		OutputSearchResult(service)
 		if count >= limit && limit != 0 {
@@ -99,4 +110,3 @@ func OutputSearchResult(searchResult LeakIXClient.SearchResult) {
 		}
 	}
 }
-