package main

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
)

var (
	flagAddress  string
	flagUsername string
	flagPassword string
)

const (
	sessionKey = "SESSION_ID_VIGOR"
)

func makeAddress(page string) string {
	u, err := url.Parse(flagAddress)
	if err != nil {
		panic(fmt.Sprintf("Invalid address (%q): %v", flagAddress, err))
	}
	u = u.JoinPath(page)
	return u.String()
}

type client struct {
	random string
	jar    *cookiejar.Jar
}

func (c *client) doRequest(ctx context.Context, body io.Reader, method, url, referer string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, method, url, body)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Referer", referer)
	req.Header.Add("Origin", makeAddress(""))

	client := &http.Client{
		Jar: c.jar,
	}

	return client.Do(req)
}

func (c *client) getSession(ctx context.Context) error {
	// Get session.
	c.random = rand.Text()[:15]
	form := url.Values{}
	form.Add("aa", base64.StdEncoding.EncodeToString([]byte(flagUsername)))
	form.Add("ab", base64.StdEncoding.EncodeToString([]byte(flagPassword)))
	form.Add("sFormAuthStr", c.random)

	res, err := c.doRequest(ctx, strings.NewReader(form.Encode()), "POST", makeAddress("cgi-bin/wlogin.cgi"), makeAddress("weblogin.htm"))
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(res.Body)

	if strings.Contains(string(body), "Authorization Error") {
		return fmt.Errorf("authorization error on wlogin.cgi")
	}

	// Check that auth works.
	res, err = c.doRequest(ctx, nil, "GET", makeAddress("cgi-bin/cgidashboard.cgi")+fmt.Sprintf("?sFormAuthStr=%s&fid=1", c.random), makeAddress(""))
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, _ = io.ReadAll(res.Body)

	if strings.Contains(string(body), "Authorization Error") {
		return fmt.Errorf("authorization error on cgidashboard.cgi")
	}

	return nil
}

func (c *client) reboot(ctx context.Context) error {
	form := url.Values{}
	form.Add("sReboot", "Current")
	form.Add("submitbnt", "Reboot Now")
	form.Add("nSch1", "0")
	form.Add("nSch2", "0")
	form.Add("nSch3", "0")
	form.Add("nSch4", "0")
	form.Add("fid", "0")
	form.Add("webchange", "0")
	form.Add("sFormAuthStr", c.random)

	res, err := c.doRequest(ctx, strings.NewReader(form.Encode()), "POST", makeAddress("cgi-bin/reboot.cgi"), makeAddress("doc/rebootsys.sht"))
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(res.Body)

	if strings.Contains(string(body), "Authorization Error") {
		return fmt.Errorf("authorization error on reboot.cgi")
	}

	return nil
}

func main() {
	flag.StringVar(&flagAddress, "address", "http://192.168.1.1/", "Address of Draytek admin interface")
	flag.StringVar(&flagUsername, "username", "admin", "Username for authenticating to the Draytek admin interface")
	flag.StringVar(&flagPassword, "password", "", "Password for authenticating to the Draytek admin interface")
	flag.Parse()

	if flagPassword == "" {
		log.Fatalf("-password must be specified")
	}

	ctx := context.Background()

	jar, _ := cookiejar.New(nil)
	c := client{
		jar: jar,
	}
	err := c.getSession(ctx)
	if err != nil {
		log.Fatalf("Could not get session: %v", err)
	}

	if err := c.reboot(ctx); err != nil {
		log.Fatalf("Reboot failed: %v", err)
	}

	log.Printf("Rebooted!")
}