forked from bettercap/bettercap
/
http_proxy_base_hosttracker.go
72 lines (61 loc) · 1.2 KB
/
http_proxy_base_hosttracker.go
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
package http_proxy
import (
"net"
"sync"
)
type Host struct {
Hostname string
Address net.IP
Resolved sync.WaitGroup
}
func NewHost(name string) *Host {
h := &Host{
Hostname: name,
Address: nil,
Resolved: sync.WaitGroup{},
}
h.Resolved.Add(1)
go func(ph *Host) {
defer ph.Resolved.Done()
if addrs, err := net.LookupIP(ph.Hostname); err == nil && len(addrs) > 0 {
ph.Address = make(net.IP, len(addrs[0]))
copy(ph.Address, addrs[0])
} else {
ph.Address = nil
}
}(h)
return h
}
type HostTracker struct {
sync.RWMutex
uhosts map[string]*Host
shosts map[string]*Host
}
func NewHostTracker() *HostTracker {
return &HostTracker{
uhosts: make(map[string]*Host),
shosts: make(map[string]*Host),
}
}
func (t *HostTracker) Track(host, stripped string) {
t.Lock()
defer t.Unlock()
t.uhosts[stripped] = NewHost(host)
t.shosts[host] = NewHost(stripped)
}
func (t *HostTracker) Unstrip(stripped string) *Host {
t.RLock()
defer t.RUnlock()
if host, found := t.uhosts[stripped]; found {
return host
}
return nil
}
func (t *HostTracker) Strip(unstripped string) *Host {
t.RLock()
defer t.RUnlock()
if host, found := t.shosts[unstripped]; found {
return host
}
return nil
}