155 lines
2.6 KiB
Go
155 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"fmt"
|
|
"io/ioutil"
|
|
"net"
|
|
"os"
|
|
"os/exec"
|
|
"os/signal"
|
|
"sort"
|
|
"strings"
|
|
"syscall"
|
|
|
|
weavenet "github.com/weaveworks/weave/net"
|
|
)
|
|
|
|
func main() {
|
|
if len(os.Args) <= 1 {
|
|
os.Exit(0)
|
|
}
|
|
|
|
var (
|
|
args = os.Args[1:]
|
|
notInExec = true
|
|
rewriteHosts = true
|
|
)
|
|
|
|
if args[0] == "-s" {
|
|
notInExec = false
|
|
rewriteHosts = false
|
|
args = args[1:]
|
|
}
|
|
|
|
if args[0] == "-h" {
|
|
rewriteHosts = false
|
|
args = args[1:]
|
|
}
|
|
|
|
if notInExec {
|
|
usr2 := make(chan os.Signal)
|
|
signal.Notify(usr2, syscall.SIGUSR2)
|
|
<-usr2
|
|
}
|
|
|
|
iface, err := weavenet.EnsureInterface("ethwe", -1)
|
|
checkErr(err)
|
|
|
|
if rewriteHosts {
|
|
updateHosts(iface)
|
|
}
|
|
|
|
binary, err := exec.LookPath(args[0])
|
|
checkErr(err)
|
|
|
|
checkErr(syscall.Exec(binary, args, os.Environ()))
|
|
}
|
|
|
|
func checkErr(err error) {
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func updateHosts(iface *net.Interface) {
|
|
addrs, err := iface.Addrs()
|
|
checkErr(err)
|
|
if len(addrs) == 0 {
|
|
return
|
|
}
|
|
hostname, err := os.Hostname()
|
|
checkErr(err)
|
|
|
|
hosts := parseHosts()
|
|
|
|
// Remove existing ips pointing to our hostname
|
|
toRemove := []string{}
|
|
for ip, addrs := range hosts {
|
|
for _, addr := range addrs {
|
|
if addr == hostname {
|
|
toRemove = append(toRemove, ip)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
for _, ip := range toRemove {
|
|
delete(hosts, ip)
|
|
}
|
|
|
|
// Add the weave ip(s)
|
|
for _, addr := range addrs {
|
|
if addr, ok := addr.(*net.IPNet); ok {
|
|
ip := addr.IP.String()
|
|
hosts[ip] = append(hosts[ip], hostname)
|
|
}
|
|
}
|
|
|
|
writeHosts(hosts)
|
|
}
|
|
|
|
func parseHosts() map[string][]string {
|
|
f, err := os.Open("/etc/hosts")
|
|
checkErr(err)
|
|
defer f.Close()
|
|
ips := map[string][]string{}
|
|
scanner := bufio.NewScanner(f)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
|
|
// Remove any comments
|
|
if i := strings.IndexByte(line, '#'); i != -1 {
|
|
line = line[:i]
|
|
}
|
|
|
|
fields := strings.Fields(line)
|
|
if len(fields) > 0 {
|
|
ips[fields[0]] = append(ips[fields[0]], fields[1:]...)
|
|
}
|
|
}
|
|
checkErr(scanner.Err())
|
|
return ips
|
|
}
|
|
|
|
func writeHosts(contents map[string][]string) {
|
|
ips := []string{}
|
|
for ip := range contents {
|
|
ips = append(ips, ip)
|
|
}
|
|
sort.Strings(ips)
|
|
|
|
buf := &bytes.Buffer{}
|
|
fmt.Fprintln(buf, "# modified by weave")
|
|
for _, ip := range ips {
|
|
if addrs := contents[ip]; len(addrs) > 0 {
|
|
fmt.Fprintf(buf, "%s\t%s\n", ip, strings.Join(uniqueStrs(addrs), " "))
|
|
}
|
|
}
|
|
checkErr(ioutil.WriteFile("/etc/hosts", buf.Bytes(), 644))
|
|
}
|
|
|
|
func uniqueStrs(s []string) []string {
|
|
m := map[string]struct{}{}
|
|
result := []string{}
|
|
for _, str := range s {
|
|
if _, ok := m[str]; !ok {
|
|
m[str] = struct{}{}
|
|
result = append(result, str)
|
|
}
|
|
}
|
|
sort.Strings(result)
|
|
return result
|
|
}
|