package main

import (
	"archive/tar"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

// Updates arrive as a tar of a signed binary cache (see lib.mkRelease):
//
//	latest
//	cache/nix-cache-info
//	cache/<hash>.narinfo
//	cache/nar/<hash>.nar[.zst|.xz]
//
// The UI only stages the files; router-update (root, a systemd unit) verifies
// the signatures. Extraction is still strict: only these names, only regular
// files, bounded in size, so a hostile bundle cannot write elsewhere.

const (
	maxBundle = 8 << 30 // bytes in all
	maxLatest = 64 << 10
)

var bundleNames = []*regexp.Regexp{
	regexp.MustCompile(`^latest$`),
	regexp.MustCompile(`^cache/nix-cache-info$`),
	regexp.MustCompile(`^cache/[0-9a-z]{32}\.narinfo$`),
	regexp.MustCompile(`^cache/nar/[0-9a-z]{20,64}\.nar(\.[a-z0-9]{1,5})?$`),
}

// unbundle extracts a bundle into dir, which it replaces.
func unbundle(r io.Reader, dir string) error {
	dir = filepath.Clean(dir)
	tmp := dir + ".tmp"
	if err := os.RemoveAll(tmp); err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Join(tmp, "cache", "nar"), 0o700); err != nil {
		return err
	}
	if err := extract(tar.NewReader(r), tmp); err != nil {
		os.RemoveAll(tmp)
		return err
	}
	for _, f := range []string{"latest", "cache/nix-cache-info"} {
		if _, err := os.Stat(filepath.Join(tmp, f)); err != nil {
			os.RemoveAll(tmp)
			return fmt.Errorf("not an update bundle: %s is missing", f)
		}
	}
	if err := os.RemoveAll(dir); err != nil {
		return err
	}
	return os.Rename(tmp, dir)
}

func extract(tr *tar.Reader, root string) error {
	var total int64
	for {
		hdr, err := tr.Next()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return fmt.Errorf("not an update bundle: %w", err)
		}
		name := path.Clean(strings.TrimPrefix(hdr.Name, "./"))
		switch hdr.Typeflag {
		case tar.TypeDir:
			// Directories are made here; the entries only have to be harmless.
			switch name {
			case ".", "cache", "cache/nar", "cache/log", "cache/realisations":
				continue
			}
			return fmt.Errorf("unexpected directory %q in the bundle", hdr.Name)
		case tar.TypeReg:
		default:
			return fmt.Errorf("unexpected entry %q in the bundle (only plain files are allowed)", hdr.Name)
		}
		ok := false
		for _, re := range bundleNames {
			ok = ok || re.MatchString(name)
		}
		if !ok {
			return fmt.Errorf("unexpected file %q in the bundle", hdr.Name)
		}
		limit := int64(maxBundle) - total
		if name == "latest" {
			limit = min(limit, maxLatest)
		}
		if hdr.Size > limit {
			return fmt.Errorf("the bundle is too large (%q)", hdr.Name)
		}
		f, err := os.OpenFile(filepath.Join(root, filepath.FromSlash(name)), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
		if err != nil {
			return fmt.Errorf("%q appears twice in the bundle", hdr.Name)
		}
		n, err := io.Copy(f, tr)
		total += n
		if cerr := f.Close(); err == nil {
			err = cerr
		}
		if err != nil {
			return err
		}
	}
}

// Release describes a version of nixos-router (/etc/router-release, and the
// latest file of an update).
type Release struct {
	Version string `json:"version"`
	Serial  int64  `json:"serial"`
	Rev     string `json:"rev"`
	Schema  int    `json:"schema"`
}

func (s System) releasePath() string {
	if s.ReleaseFile != "" {
		return s.ReleaseFile
	}
	return "/etc/router-release"
}

func (s System) installedRelease() Release {
	var r Release
	if data, err := os.ReadFile(s.releasePath()); err == nil {
		json.Unmarshal(data, &r)
	}
	return r
}

func (s System) updateDir() string     { return filepath.Join(s.StateDir, "update") }
func (s System) updateURLPath() string { return filepath.Join(s.StateDir, "webui", "update.json") }

func (s System) updateURL() string {
	var v struct {
		URL string `json:"url"`
	}
	if data, err := os.ReadFile(s.updateURLPath()); err == nil {
		json.Unmarshal(data, &v)
	}
	return v.URL
}

func (s System) saveUpdateURL(u string) error {
	data, _ := json.Marshal(struct {
		URL string `json:"url"`
	}{u})
	if err := os.MkdirAll(filepath.Dir(s.updateURLPath()), 0o700); err != nil {
		return err
	}
	return writeFileAtomic(s.updateURLPath(), data, 0o600)
}

// parseUpdateURL accepts the base of an update server: http or https, no
// credentials, no query. The data is signed, so plain HTTP is no risk to the
// router; it only lets an attacker withhold updates.
func parseUpdateURL(v string) (string, error) {
	v = strings.TrimRight(strings.TrimSpace(v), "/")
	u, err := url.Parse(v)
	if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil ||
		u.RawQuery != "" || u.Fragment != "" {
		return "", errors.New("an http:// or https:// address without credentials")
	}
	return v, nil
}

// Available is what the update server offers.
type Available struct {
	Release
	Newer bool
}

func (s System) checkUpdate(base string) (Available, error) {
	client := &http.Client{Timeout: 15 * time.Second}
	resp, err := client.Get(base + "/latest")
	if err != nil {
		return Available{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Available{}, fmt.Errorf("the server answered %s", resp.Status)
	}
	var a Available
	if err := json.NewDecoder(io.LimitReader(resp.Body, maxLatest)).Decode(&a); err != nil || a.Version == "" {
		return Available{}, errors.New("the server's latest file is not an update description")
	}
	a.Newer = a.Serial > s.installedRelease().Serial
	return a, nil
}

// request asks the router-update unit to install; it reads this file.
func (s System) requestUpdate(mode, url string, force bool) error {
	data, _ := json.Marshal(map[string]any{"mode": mode, "url": url, "force": force})
	if err := os.MkdirAll(s.updateDir(), 0o700); err != nil {
		return err
	}
	if err := writeFileAtomic(filepath.Join(s.updateDir(), "request.json"), data, 0o600); err != nil {
		return err
	}
	return startUnit("router-update.service")
}

// startUnit is replaced in tests.
var startUnit = func(unit string) error {
	_, err := command("systemctl", "start", "--no-block", unit)
	return err
}

// UpdateStatus is the state of a one-shot unit (router-update, router-upgrade)
// and what its last run logged.
type UpdateStatus struct {
	Running bool
	Result  string // of the last run; "" if none
	Log     string
}

func unitRun(unit string) UpdateStatus {
	st := UpdateStatus{}
	st.Running = unitProperty(unit, "ActiveState") == "activating"
	if inv := unitProperty(unit, "InvocationID"); inv != "" {
		if !st.Running {
			st.Result = unitProperty(unit, "Result")
		}
		st.Log, _ = command("journalctl", "--no-pager", "-o", "cat", "-n", "200", "_SYSTEMD_INVOCATION_ID="+inv)
	}
	return st
}

func (s System) UpdateStatus() UpdateStatus { return unitRun("router-update.service") }

type updateData struct {
	Installed Release
	Status    UpdateStatus
	URL       string
	Offer     *Available
	OfferErr  string
	Pending   bool // an update waits for confirmation
}

func (a *app) updateData() updateData {
	st := a.sys.Status()
	return updateData{
		Installed: a.sys.installedRelease(),
		Status:    a.sys.UpdateStatus(),
		URL:       a.sys.updateURL(),
		Pending:   st.RollbackIn > 0,
	}
}

func (a *app) updatePage(w http.ResponseWriter, r *http.Request, sess session) {
	d := a.updateData()
	p := page{Title: "Update", CSRF: sess.csrf, LoggedIn: true, Data: d, Flash: r.URL.Query().Get("msg")}
	if d.Status.Running {
		p.Refresh = 3
	}
	a.render(w, "update", p)
}

func (a *app) updateBack(w http.ResponseWriter, r *http.Request, msg string) {
	http.Redirect(w, r, "/update?msg="+url.QueryEscape(msg), http.StatusSeeOther)
}

func (a *app) busy() bool {
	st := a.sys.Status()
	return st.Applying || st.ConfirmRunning || st.RollbackIn > 0 || a.sys.UpdateStatus().Running || a.sys.UpgradeStatus().Running
}

func (a *app) updateSaveURL(w http.ResponseWriter, r *http.Request, _ session) {
	v := strings.TrimSpace(r.PostFormValue("url"))
	if v != "" {
		var err error
		if v, err = parseUpdateURL(v); err != nil {
			a.updateBack(w, r, "Not saved: "+err.Error())
			return
		}
	}
	if err := a.sys.saveUpdateURL(v); err != nil {
		a.updateBack(w, r, "Not saved: "+err.Error())
		return
	}
	a.updateBack(w, r, "Saved.")
}

func (a *app) updateCheck(w http.ResponseWriter, r *http.Request, sess session) {
	d := a.updateData()
	p := page{Title: "Update", CSRF: sess.csrf, LoggedIn: true}
	if d.URL == "" {
		a.updateBack(w, r, "Set the update server first.")
		return
	}
	offer, err := a.sys.checkUpdate(d.URL)
	if err != nil {
		d.OfferErr = err.Error()
	} else {
		d.Offer = &offer
	}
	p.Data = d
	a.render(w, "update", p)
}

func (a *app) updateInstall(w http.ResponseWriter, r *http.Request, _ session) {
	if a.busy() {
		a.updateBack(w, r, "Another change is in progress or waiting for confirmation.")
		return
	}
	u := a.sys.updateURL()
	if u == "" {
		a.updateBack(w, r, "Set the update server first.")
		return
	}
	if err := a.sys.requestUpdate("url", u, r.PostFormValue("force") == "on"); err != nil {
		a.updateBack(w, r, "Could not start the update: "+err.Error())
		return
	}
	a.updateBack(w, r, "Updating.")
}

// updateUpload takes a bundle as multipart/form-data and streams it into the
// staging directory; a bundle is hundreds of megabytes. The CSRF token is in
// the query string, since the body must not be read before it is checked.
func (a *app) updateUpload(w http.ResponseWriter, r *http.Request) {
	sess, ok := a.session(r)
	if !ok {
		http.Redirect(w, r, "/login", http.StatusSeeOther)
		return
	}
	if !sameOrigin(r) || r.URL.Query().Get("csrf") != sess.csrf {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}
	if a.busy() {
		a.updateBack(w, r, "Another change is in progress or waiting for confirmation.")
		return
	}
	mr, err := r.MultipartReader()
	if err != nil {
		a.updateBack(w, r, "Not an upload: "+err.Error())
		return
	}
	force := false
	for {
		part, err := mr.NextPart()
		if errors.Is(err, io.EOF) {
			a.updateBack(w, r, "No file was uploaded.")
			return
		}
		if err != nil {
			a.updateBack(w, r, "Upload failed: "+err.Error())
			return
		}
		if part.FormName() == "force" {
			force = true
			continue
		}
		if part.FormName() == "bundle" && part.FileName() != "" {
			a.stageBundle(w, r, part, force)
			return
		}
	}
}

func (a *app) stageBundle(w http.ResponseWriter, r *http.Request, part *multipart.Part, force bool) {
	if err := os.MkdirAll(a.sys.updateDir(), 0o700); err != nil {
		a.updateBack(w, r, "Could not stage the bundle: "+err.Error())
		return
	}
	if err := unbundle(part, filepath.Join(a.sys.updateDir(), "staged")); err != nil {
		a.updateBack(w, r, "Rejected: "+err.Error())
		return
	}
	if err := a.sys.requestUpdate("staged", "", force); err != nil {
		a.updateBack(w, r, "Could not start the update: "+err.Error())
		return
	}
	a.updateBack(w, r, "Bundle received; verifying and installing.")
}