forked from remind101/empire
/
awsutil.go
109 lines (90 loc) · 2.05 KB
/
awsutil.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
package awsutil
import (
"bytes"
"encoding/json"
"fmt"
"io"
"io/ioutil"
"net/http"
"strings"
)
// Request represents an expected AWS API Operation.
type Request struct {
RequestURI string
Operation string
Body string
}
func (r *Request) String() string {
body := formatBody(strings.NewReader(r.Body))
return fmt.Sprintf("RequestURI: %s\nOperation: %s\nBody: %s", r.RequestURI, r.Operation, body)
}
// Response represents a predefined response.
type Response struct {
StatusCode int
Body string
}
// Cycle represents a request-response cycle.
type Cycle struct {
Request Request
Response Response
}
// Handler is an http.Handler that will play back cycles.
type Handler struct {
cycles []Cycle
}
// NewHandler returns a new Handler instance.
func NewHandler(c []Cycle) *Handler {
return &Handler{cycles: c}
}
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if len(h.cycles) == 0 {
fmt.Println("No cycles remaining to replay.")
w.WriteHeader(404)
return
}
b, err := ioutil.ReadAll(r.Body)
if err != nil {
panic(err)
}
cycle := h.cycles[0]
match := Request{
RequestURI: r.URL.RequestURI(),
Operation: r.Header.Get("X-Amz-Target"),
Body: string(b),
}
if cycle.Request.Body == "ignore" {
match.Body = cycle.Request.Body
}
if cycle.Request.String() == match.String() {
w.WriteHeader(cycle.Response.StatusCode)
io.WriteString(w, cycle.Response.Body)
} else {
fmt.Println("Request does not match next cycle.")
fmt.Println(cycle.Request.String())
fmt.Println(match.String())
w.WriteHeader(404)
}
h.cycles = h.cycles[1:]
}
func formatBody(r io.Reader) string {
b, err := ioutil.ReadAll(r)
if err != nil {
panic(err)
}
s, err := formatJSON(bytes.NewReader(b))
if err == nil {
return s
}
return string(b)
}
func formatJSON(r io.Reader) (string, error) {
var body map[string]interface{}
if err := json.NewDecoder(r).Decode(&body); err != nil {
return "", err
}
raw, err := json.MarshalIndent(&body, "", " ")
if err != nil {
return "", err
}
return string(raw), nil
}