random-redirect/main.go

98 lines
2.0 KiB
Go

package main
import (
"bytes"
"flag"
"fmt"
"io"
"io/ioutil"
"log"
"math/rand"
"net/http"
"strconv"
"strings"
"time"
"golang.org/x/time/rate"
)
func main() {
rand.Seed(time.Now().UnixNano())
var port int
var conf string
var proxy bool
flag.IntVar(&port, "p", 8080, "port to listen on")
flag.StringVar(&conf, "c", "/dev/null", "line delimited file to read urls from")
flag.BoolVar(&proxy, "proxy", false, "proxy content rather than redirect, so refresh works")
flag.Parse()
b, err := ioutil.ReadFile(conf)
if err != nil {
panic(err)
}
links := bytes.Split(b, []byte{'\n'})
log.Print(port)
server := Server{
limiter: rate.NewLimiter(1, 1),
links: links,
proxy: proxy,
}
if err := http.ListenAndServe(fmt.Sprintf(":%d", port), server); err != nil {
panic(err)
}
}
type Server struct {
limiter *rate.Limiter
proxy bool
links [][]byte
}
func (server Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
s := r.URL.Query().Get("n")
if s == "" {
s = "1"
}
n, err := strconv.Atoi(s)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
links := make([]string, n)
for i := range links {
link := ""
for len(link) == 0 || strings.Contains(fmt.Sprint(links), link) {
link = string(server.links[rand.Intn(len(server.links))])
}
links[i] = link
}
for _, link := range links {
if strings.HasPrefix(link, "http") {
server.serveHTTP(w, r, link)
} else if strings.HasPrefix(link, "literal://") {
w.Write([]byte(strings.TrimPrefix(link, "literal://")))
}
w.Write([]byte("\n"))
}
}
func (server Server) serveHTTP(w http.ResponseWriter, r *http.Request, link string) {
if server.proxy {
server.limiter.Wait(r.Context())
resp, err := http.Get(link)
if err != nil {
http.Error(w, err.Error(), http.StatusBadGateway)
}
for k, v := range resp.Header {
for _, v2 := range v {
w.Header().Add(k, v2)
}
}
w.WriteHeader(resp.StatusCode)
io.Copy(w, resp.Body)
} else {
http.Redirect(w, r, link, http.StatusTemporaryRedirect)
}
}