interfare/dtd_download.go

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/BurntSushi/toml"
)

const DEFAULT_NRDP_BASE_URL = "https://opendata.nationalrail.co.uk"
const DEFAULT_USER_AGENT = "interfare/0.1 (interfare@eta.st)"

type Config struct {
	BaseURL      string `toml:"base_url"`
	NrdpUsername string `toml:"nrdp_username"`
	NrdpPassword string `toml:"nrdp_password"`
	UserAgent    string `toml:"user_agent"`
}

func loadConfig(path string) (Config, error) {
	var ret Config
	data, err := os.ReadFile(path)
	if err != nil {
		return ret, fmt.Errorf("failed to read config file %s: %w", path, err)
	}
	_, err = toml.Decode(string(data), &ret)
	if err != nil {
		return ret, fmt.Errorf("failed to decode config file %s: %w", path, err)
	}
	if ret.BaseURL == "" {
		ret.BaseURL = DEFAULT_NRDP_BASE_URL
	}
	if ret.UserAgent == "" {
		ret.UserAgent = DEFAULT_USER_AGENT
	}
	if ret.NrdpUsername == "" || ret.NrdpPassword == "" {
		return ret, fmt.Errorf("must specify `nrdp_username` and `nrdp_password` in config file")
	}
	return ret, nil
}

type NrdpClient struct {
	cfg         Config
	apiToken    string
	tokenExpiry time.Time
}

func NewNrdpClient(cfg Config) NrdpClient {
	var ret NrdpClient
	ret.cfg = cfg
	return ret
}

// EnsureToken checks that the NrdpClient's API token has not expired, and calls /authenticate to get a new one
// if it has.
func (cli *NrdpClient) EnsureToken() error {
	if cli.tokenExpiry.After(time.Now()) {
		// don't need to do this
		return nil
	}
	uri := fmt.Sprintf("%s/authenticate", cli.cfg.BaseURL)
	form := url.Values{}
	form.Add("username", cli.cfg.NrdpUsername)
	form.Add("password", cli.cfg.NrdpPassword)
	req, err := http.NewRequest("POST", uri, strings.NewReader(form.Encode()))
	if err != nil {
		return fmt.Errorf("failed to build token request: %w", err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("User-Agent", cli.cfg.UserAgent)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("token request failed: %w", err)
	}
	body := make(map[string]interface{})
	err = json.NewDecoder(resp.Body).Decode(&body)
	if err != nil {
		return fmt.Errorf("failed to decode token response json: %w", err)
	}
	token, ok := body["token"]
	if !ok {
		return fmt.Errorf("token response did not contain token; body %#v", body)
	}
	tok, ok := token.(string)
	if !ok {
		return fmt.Errorf("token response was not string; body %#v", body)
	}
	cli.apiToken = tok
	cli.tokenExpiry = time.Now().Add(1 * time.Hour)
	return nil
}

// FetchFile obtains the named file (one of "fares", "routeing", or "timetable") from NRDP
// and returns an interface to read it.
func (cli *NrdpClient) FetchFile(name string) (io.ReadCloser, error) {
	if err := cli.EnsureToken(); err != nil {
		return nil, fmt.Errorf("failed to ensure token: %w", err)
	}
	var version string
	switch name {
	case "routeing":
	case "fares":
		version = "2.0"
	case "timetable":
		version = "3.0"
	default:
		panic(fmt.Sprintf("invalid argument %s to FetchFile", name))
	}
	uri := fmt.Sprintf("%s/api/staticfeeds/%s/%s", cli.cfg.BaseURL, version, name)
	req, err := http.NewRequest("GET", uri, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to build file request: %w", err)
	}
	req.Header.Set("X-Auth-Token", cli.apiToken)
	req.Header.Set("User-Agent", cli.cfg.UserAgent)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("file http get failed: %w", err)
	}
	return resp.Body, nil
}

var configPath = flag.String("config", "./config.toml", "Path to the configuration file to use.")

func main() {
	flag.Parse()
	cfg, err := loadConfig(*configPath)
	if err != nil {
		log.Fatalf("failed to load config: %s", err)
	}
	if _, err := os.Stat("./fares.zip"); errors.Is(err, os.ErrNotExist) {
		client := NewNrdpClient(cfg)
		err = client.EnsureToken()
		if err != nil {
			log.Fatalf("failed to get token: %s", err)
		}
		out, err := os.Create("./fares.zip")
		if err != nil {
			log.Fatalf("failed to create ./fares.zip: %s", err)
		}
		faresData, err := client.FetchFile("fares")
		if err != nil {
			log.Fatalf("failed to fetch fares data: %s", err)
		}
		written, err := io.Copy(out, faresData)
		if err != nil {
			log.Fatalf("failed to write fares data: %s", err)
		}
		log.Printf("wrote %d bytes of fares data", written)
	} else {
		log.Printf("using pre-existing fares.zip")
	}
	err = ParseFaresData("./fares.zip")
	if err != nil {
		log.Fatalf("failed to parse fares data: %s", err)
	}
}