overthewire - natas/request.go

     1  package main
     2  
     3  import (
     4  	"io"
     5  	"mime/multipart"
     6  	"net/http"
     7  	"net/url"
     8  
     9  	"vimagination.zapto.org/errors"
    10  	"vimagination.zapto.org/memio"
    11  )
    12  
    13  var buf memio.Buffer
    14  
    15  type Grabber interface {
    16  	Grab(http.Request) string
    17  }
    18  
    19  type Grab struct{}
    20  
    21  var grab Grab
    22  
    23  func (g Grab) Grab(r http.Request) string {
    24  	buf = buf[:0]
    25  	resp, err := http.DefaultClient.Do(&r)
    26  	e(err)
    27  	io.Copy(&buf, resp.Body)
    28  	resp.Body.Close()
    29  	return string(buf)
    30  }
    31  
    32  type Text struct {
    33  	Text string
    34  }
    35  
    36  func (t Text) Grab(http.Request) string {
    37  	return t.Text
    38  }
    39  
    40  type Path struct {
    41  	Grabber
    42  	Path Grabber
    43  }
    44  
    45  func (p Path) Grab(r http.Request) string {
    46  	oldPath := r.URL.Path
    47  	r.URL.Path = p.Path.Grab(r)
    48  	s := p.Grabber.Grab(r)
    49  	r.URL.Path = oldPath
    50  	return s
    51  }
    52  
    53  type Host struct {
    54  	Grabber
    55  	Host Grabber
    56  }
    57  
    58  func (h Host) Grab(r http.Request) string {
    59  	oldHost := r.URL.Host
    60  	r.URL.Host = h.Host.Grab(r)
    61  	s := h.Grabber.Grab(r)
    62  	r.URL.Host = oldHost
    63  	return s
    64  }
    65  
    66  type Headers struct {
    67  	Grabber
    68  	SetData
    69  }
    70  
    71  func (h Headers) Grab(r http.Request) string {
    72  	newHeaders := make(http.Header)
    73  	for key, value := range r.Header {
    74  		newHeaders[key] = value
    75  	}
    76  	h.Set(r, newHeaders)
    77  	r.Header = newHeaders
    78  	return h.Grabber.Grab(r)
    79  }
    80  
    81  type Post struct {
    82  	Grabber
    83  	SetData
    84  	File *File
    85  }
    86  
    87  type File struct {
    88  	Field, Name Grabber
    89  	memio.Buffer
    90  }
    91  
    92  type MPSetter struct {
    93  	*multipart.Writer
    94  }
    95  
    96  func (m MPSetter) Set(key, value string) {
    97  	e(m.WriteField(key, value))
    98  }
    99  
   100  func (p Post) Grab(r http.Request) string {
   101  	var b memio.Buffer
   102  	m := multipart.NewWriter(&b)
   103  	p.Set(r, MPSetter{m})
   104  	if p.File != nil {
   105  		rw, err := m.CreateFormFile(p.File.Field.Grab(r), p.File.Name.Grab(r))
   106  		e(err)
   107  		io.Copy(rw, p.File)
   108  	}
   109  	e(m.Close())
   110  	r.Body = &b
   111  	r.Method = http.MethodPost
   112  	r.ContentLength = int64(len(b))
   113  	return Headers{p.Grabber, map[string]Grabber{"Content-Type": Text{m.FormDataContentType()}}}.Grab(r)
   114  }
   115  
   116  type Get struct {
   117  	Grabber
   118  	SetData
   119  }
   120  
   121  func (g Get) Grab(r http.Request) string {
   122  	values := make(url.Values, len(g.SetData))
   123  	g.Set(r, values)
   124  	oldQuery := r.URL.RawQuery
   125  	r.URL.RawQuery = values.Encode()
   126  	s := g.Grabber.Grab(r)
   127  	r.URL.RawPath = oldQuery
   128  	return s
   129  }
   130  
   131  type SetData map[string]Grabber
   132  
   133  type SetDataGrabber interface {
   134  	Grabber
   135  	SetKey(string, Grabber)
   136  }
   137  
   138  type Setter interface {
   139  	Set(string, string)
   140  }
   141  
   142  func (sd SetData) Set(r http.Request, s Setter) {
   143  	for key, g := range sd {
   144  		s.Set(key, g.Grab(r))
   145  	}
   146  }
   147  
   148  func (sd SetData) SetKey(Key string, Value Grabber) {
   149  	sd[Key] = Value
   150  }
   151  
   152  type Cookie struct {
   153  	Name string
   154  }
   155  
   156  func (c Cookie) Grab(r http.Request) string {
   157  	r.Method = http.MethodHead
   158  	resp, err := http.DefaultClient.Do(&r)
   159  	e(err)
   160  	cookies := resp.Cookies()
   161  	for _, cookie := range cookies {
   162  		if cookie.Name == c.Name {
   163  			return cookie.Value
   164  		}
   165  	}
   166  	panic(errors.Error("could not find cookie"))
   167  }
   168  
   169  type GetHeader struct {
   170  	Header string
   171  }
   172  
   173  func (g GetHeader) Grab(r http.Request) string {
   174  	resp, err := http.DefaultClient.Do(&r)
   175  	e(err)
   176  	e(resp.Body.Close())
   177  	return resp.Header.Get(g.Header)
   178  }
   179  
   180  type LoadAll []Grabber
   181  
   182  func (l LoadAll) Grab(r http.Request) string {
   183  	for _, g := range l[:len(l)-1] {
   184  		g.Grab(r)
   185  	}
   186  	return l[len(l)-1].Grab(r)
   187  }
   188  
   189  type PostBody struct {
   190  	Grabber
   191  	memio.Buffer
   192  }
   193  
   194  func (p PostBody) Grab(r http.Request) string {
   195  	r.Body = &p.Buffer
   196  	r.Method = http.MethodPost
   197  	r.ContentLength = int64(len(p.Buffer))
   198  	return Headers{p.Grabber, map[string]Grabber{"Content-Type": Text{"application/x-www-form-urlencoded"}}}.Grab(r)
   199  }
   200  
   201  type Query struct {
   202  	Grabber
   203  	RawQuery Grabber
   204  }
   205  
   206  func (q Query) Grab(r http.Request) string {
   207  	oldQuery := r.URL.RawQuery
   208  	r.URL.RawQuery = q.RawQuery.Grab(r)
   209  	str := q.Grabber.Grab(r)
   210  	r.URL.RawPath = oldQuery
   211  	return str
   212  }