wannabe/main.go at master · trco/wannabe · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/AdguardTeam/gomitmproxy"
	"github.com/AdguardTeam/gomitmproxy/mitm"
	"github.com/trco/wannabe/internal/config"
	"github.com/trco/wannabe/internal/handler"
	"github.com/trco/wannabe/internal/storage"
)

func main() {
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatalf("fatal error starting app: %v", err)
	}

	storageProvider, err := storage.ProviderFactory(cfg.StorageProvider)
	if err != nil {
		log.Fatalf("fatal error starting app: %v", err)
	}

	mitmConfig, err := config.LoadMitmConfig()
	if err != nil {
		log.Fatalf("fatal error starting app: %v", err)
	}

	go startWannabeApiServer(cfg, storageProvider)
	startWannabeProxyServer(cfg, mitmConfig, storageProvider)
}

func startWannabeApiServer(cfg config.Config, storageProvider storage.Provider) {
	http.HandleFunc("/wannabe/api/records/{hash}", handler.Records(cfg, storageProvider))
	http.HandleFunc("/wannabe/api/records", handler.Records(cfg, storageProvider))
	http.HandleFunc("/wannabe/api/regenerate", handler.Regenerate(cfg, storageProvider))

	fmt.Println("API server start listening to [::]:6790")
	err := http.ListenAndServe(":6790", nil)
	if err != nil {
		log.Fatal(err)
	}
}

func startWannabeProxyServer(cfg config.Config, mitmConfig *mitm.Config, storageProvider storage.Provider) {
	wannabeProxy := initWannabeProxy(cfg, mitmConfig, storageProvider)

	err := wannabeProxy.Start()
	if err != nil {
		log.Fatal(err)
	}

	signalChannel := make(chan os.Signal, 1)
	signal.Notify(signalChannel, syscall.SIGINT, syscall.SIGTERM)
	<-signalChannel

	wannabeProxy.Close()
}

func initWannabeProxy(cfg config.Config, mitmConfig *mitm.Config, storageProvider storage.Provider) *gomitmproxy.Proxy {
	wannabeProxy := gomitmproxy.NewProxy(gomitmproxy.Config{
		ListenAddr: &net.TCPAddr{
			IP:   net.IPv4(0, 0, 0, 0),
			Port: 6789,
		},
		MITMConfig: mitmConfig,
		OnRequest:  handler.Request(cfg, storageProvider),
		OnResponse: handler.Response(cfg, storageProvider),
	})

	return wannabeProxy
}