wswg/main.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"hash/fnv"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"sync"
	"time"
)

var version = "dev"

type App struct {
	mu       sync.Mutex
	dir      string
	settings Settings
	engine   Engine
	noTunnel bool
}

func main() {
	if e := run(); e != nil {
		fmt.Fprintln(os.Stderr, "wswg:", e)
		os.Exit(1)
	}
}
func run() error {
	enableTerminal()
	var auto, daemon, tui, gui, ver, noTunnel bool
	var config, dir string
	flag.BoolVar(&auto, "a", false, "connect last profile")
	flag.BoolVar(&auto, "auto", false, "connect last profile")
	flag.BoolVar(&daemon, "d", false, "run in background, requires -a or -c")
	flag.BoolVar(&daemon, "daemon", false, "run in background")
	flag.BoolVar(&tui, "t", false, "force terminal UI")
	flag.BoolVar(&tui, "tui", false, "force terminal UI")
	flag.BoolVar(&gui, "g", false, "force browser GUI")
	flag.BoolVar(&gui, "gui", false, "force browser GUI")
	flag.StringVar(&config, "c", "", "profile number or name")
	flag.StringVar(&config, "config", "", "profile number or name")
	flag.StringVar(&dir, "data-dir", "", "settings directory")
	flag.BoolVar(&ver, "version", false, "show version")
	flag.BoolVar(&noTunnel, "no-wstunnel", false, "disable wstunnel for this -a / -c connection only")
	flag.Parse()
	if ver {
		fmt.Println("wswg", version, runtime.GOOS, runtime.GOARCH)
		return nil
	}
	if tui && gui {
		return errors.New("-t and -g cannot be combined")
	}
	if flag.NArg() > 0 {
		return errors.New("unexpected positional arguments")
	}
	if dir == "" {
		base, e := os.UserConfigDir()
		if e != nil {
			return e
		}
		dir = filepath.Join(base, "wswg")
	}
	var e error
	dir, e = filepath.Abs(dir)
	if e != nil {
		return e
	}
	if e = os.MkdirAll(dir, 0700); e != nil {
		return e
	}
	if e = os.Chmod(dir, 0700); e != nil {
		return e
	}
	s, e := loadSettings(dir)
	if e != nil {
		return e
	}
	if noTunnel && !auto && config == "" {
		return errors.New("--no-wstunnel requires -a or -c")
	}
	selected := config
	if selected == "" && auto {
		selected = s.Last
		if selected == "" {
			return errors.New("no previous profile; choose -c first")
		}
	}
	if daemon {
		if selected == "" {
			return errors.New("-d requires -a or -c")
		}
		if _, e = s.find(selected); e != nil {
			return e
		}
		exe, e := os.Executable()
		if e != nil {
			return e
		}
		log, e := os.OpenFile(filepath.Join(dir, "daemon.log"), os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
		if e != nil {
			return e
		}
		defer log.Close()
		args := []string{"--data-dir", dir, "-c", selected}
		if noTunnel {
			args = append(args, "--no-wstunnel")
		}
		c := exec.Command(exe, args...)
		ready := filepath.Join(dir, ".ready-"+newID())
		defer os.Remove(ready)
		c.Env = append(os.Environ(), "WSWG_DAEMON=1", "WSWG_READY="+ready)
		c.Stdout = log
		c.Stderr = log
		background(c)
		if e = c.Start(); e != nil {
			return e
		}
		wait := make(chan error, 1)
		go func() { wait <- c.Wait() }()
		deadline := time.NewTimer(30 * time.Second)
		defer deadline.Stop()
		ticker := time.NewTicker(100 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case err := <-wait:
				return fmt.Errorf("daemon exited: %v; see %s", err, log.Name())
			case <-deadline.C:
				_ = c.Process.Kill()
				<-wait
				return fmt.Errorf("daemon startup timed out; see %s", log.Name())
			case <-ticker.C:
				if _, err := os.Stat(ready); err == nil {
					fmt.Println("Daemon connected, PID", c.Process.Pid, "log:", log.Name())
					return nil
				}
			}
		}
	}
	h := fnv.New32a()
	h.Write([]byte(dir))
	lock, e := net.Listen("tcp4", fmt.Sprintf("127.0.0.1:%d", 30000+h.Sum32()%20000))
	if e != nil {
		return errors.New("another wswg instance is running for this data directory, or its lock port is busy")
	}
	defer lock.Close()
	a := &App{dir: dir, settings: s, noTunnel: noTunnel}
	a.engine.dir = dir
	defer a.engine.stop()
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt, terminationSignal())
	defer signal.Stop(signals)
	done := make(chan struct{})
	defer close(done)
	go func() {
		select {
		case <-signals:
			a.engine.stop()
			os.Exit(0)
		case <-done:
		}
	}()
	if selected != "" {
		p, e := s.find(selected)
		if e != nil {
			return e
		}
		if e = a.connect(p.ID); e != nil {
			return e
		}
		if os.Getenv("WSWG_DAEMON") == "1" {
			if ready := os.Getenv("WSWG_READY"); ready != "" {
				if e = os.WriteFile(ready, []byte("ready"), 0600); e != nil {
					return e
				}
			}
			for {
				time.Sleep(time.Second)
				if s := a.engine.status(); !s.Connected {
					return fmt.Errorf("daemon disconnected: %s", s.Error)
				}
			}
		}
		for {
			if s := a.engine.status(); !s.Connected {
				return fmt.Errorf("disconnected: %s", s.Error)
			}
			fmt.Printf("\r\033[K%s", a.engine.status().Line())
			time.Sleep(time.Second)
		}
	}
	if !tui && !gui {
		gui = runtime.GOOS == "windows" || os.Getenv("DISPLAY") != "" || os.Getenv("WAYLAND_DISPLAY") != ""
	}
	if gui {
		if s.Autorun && s.Last != "" {
			if e = a.connect(s.Last); e != nil {
				fmt.Fprintln(os.Stderr, e)
			}
		}
		return a.gui()
	}
	return a.tui()
}
func (a *App) connect(id string) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	p, e := a.settings.find(id)
	if e != nil {
		return e
	}
	if a.noTunnel {
		p.Tunnel.Enabled = false
	}
	if e = a.engine.start(p); e != nil {
		return e
	}
	old := a.settings.Last
	a.settings.Last = p.ID
	if e = saveSettings(a.dir, a.settings); e != nil {
		a.settings.Last = old
		a.engine.stop()
	}
	return e
}
func (a *App) save(p Profile) error {
	if e := validateProfile(p); e != nil {
		return e
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.engine.activeID() == p.ID && p.ID != "" {
		return errors.New("disconnect before editing this profile")
	}
	for _, x := range a.settings.Profiles {
		if x.Name == p.Name && x.ID != p.ID {
			return errors.New("name already exists")
		}
	}
	old := append([]Profile(nil), a.settings.Profiles...)
	found := false
	for i, x := range a.settings.Profiles {
		if x.ID == p.ID && p.ID != "" {
			a.settings.Profiles[i] = p
			found = true
		}
	}
	if !found {
		p.ID = newID()
		a.settings.Profiles = append(a.settings.Profiles, p)
	}
	if e := saveSettings(a.dir, a.settings); e != nil {
		a.settings.Profiles = old
		return e
	}
	return nil
}
func (a *App) remove(id string) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.engine.activeID() == id {
		return errors.New("disconnect before deleting")
	}
	old := a.settings
	next := a.settings
	next.Profiles = []Profile{}
	for _, p := range old.Profiles {
		if p.ID != id {
			next.Profiles = append(next.Profiles, p)
		}
	}
	if next.Last == id {
		next.Last = ""
	}
	if e := saveSettings(a.dir, next); e != nil {
		return e
	}
	a.settings = next
	return nil
}