xonstat/util/xs_interceptor/xs_interceptor.go

   1 package main\r
   2 \r
   3 import "database/sql"\r
   4 import "flag"\r
   5 import "fmt"\r
   6 import "html/template"\r
   7 import "net/http"\r
   8 import "os"\r
   9 import "strings"\r
  10 import _ "github.com/mattn/go-sqlite3"\r
  11 \r
  12 // HTML templates\r
  13 var templates = template.Must(template.ParseFiles("templates/landing.html"))\r
  14 \r
  15 func main() {\r
  16         port := flag.Int("port", 6543, "Default port on which to accept requests")\r
  17         url := flag.String("url", "http://localhost:6543/stats/submit", "URL to send POST requests against")\r
  18         flag.Usage = usage\r
  19         flag.Parse()\r
  20 \r
  21         if len(flag.Args()) < 1 {\r
  22                 fmt.Println("Insufficient arguments: need a <command> to run. Exiting...")\r
  23                 os.Exit(1)\r
  24         }\r
  25 \r
  26         command := flag.Args()[0]\r
  27         switch {\r
  28         case command == "drop":\r
  29                 drop_db()\r
  30         case command == "create":\r
  31                 create_db()\r
  32         case command == "serve":\r
  33                 serve(*port)\r
  34         case command == "resubmit":\r
  35                 resubmit(*url)\r
  36         case command == "list":\r
  37                 list()\r
  38         default:\r
  39                 fmt.Println("Unknown command! Exiting...")\r
  40                 os.Exit(1)\r
  41         }\r
  42 }\r
  43 \r
  44 // override the default Usage function to show the different "commands"\r
  45 // that are in the switch statement in main()\r
  46 func usage() {\r
  47         fmt.Fprintf(os.Stderr, "Usage of xs_interceptor:\n")\r
  48         fmt.Fprintf(os.Stderr, "    xs_interceptor [options] <command>\n\n")\r
  49         fmt.Fprintf(os.Stderr, "Where <command> is one of the following:\n")\r
  50         fmt.Fprintf(os.Stderr, "    create   - create the requests db (sqlite3 db file)\n")\r
  51         fmt.Fprintf(os.Stderr, "    drop     - remove the requests db\n")\r
  52         fmt.Fprintf(os.Stderr, "    list     - lists the requests in the db\n")\r
  53         fmt.Fprintf(os.Stderr, "    serve    - listens for stats requests, storing them if found\n")\r
  54         fmt.Fprintf(os.Stderr, "    resubmit - resubmits the requests to another URL\n\n")\r
  55         fmt.Fprintf(os.Stderr, "Where [options] is one or more of the following:\n")\r
  56         fmt.Fprintf(os.Stderr, "    -port    - port number (int) to listen on for 'serve' command\n")\r
  57         fmt.Fprintf(os.Stderr, "    -url     - url (string) to submit requests\n\n")\r
  58 }\r
  59 \r
  60 // removes the requests database. it is just a file, so this is really easy.\r
  61 func drop_db() {\r
  62         err := os.Remove("middleman.db")\r
  63 \r
  64         if err != nil {\r
  65                 fmt.Println("Error dropping the database middleman.db. Exiting...")\r
  66                 os.Exit(1)\r
  67         } else {\r
  68                 fmt.Println("Dropped middleman.db successfully!")\r
  69                 os.Exit(0)\r
  70         }\r
  71 }\r
  72 \r
  73 // creates the sqlite database. it's a hard-coded name because I don't see\r
  74 // a need to change db names for this purpose.\r
  75 func create_db() {\r
  76         db, err := sql.Open("sqlite3", "./middleman.db")\r
  77         defer db.Close()\r
  78 \r
  79         if err != nil {\r
  80                 fmt.Println("Error creating the database middleman.db. Exiting...")\r
  81                 fmt.Println(err)\r
  82                 os.Exit(1)\r
  83         } else {\r
  84                 fmt.Println("Created middleman.db successfully!")\r
  85         }\r
  86 \r
  87         _, err = db.Exec(`\r
  88      CREATE TABLE requests (\r
  89         request_id INTEGER PRIMARY KEY ASC, \r
  90         blind_id_header TEXT, \r
  91         ip_addr VARCHAR(32), \r
  92         body TEXT, \r
  93         bodylength int \r
  94      );\r
  95   `)\r
  96 \r
  97         if err != nil {\r
  98                 fmt.Println("Error creating the table 'requests' in middleman.db. Exiting...")\r
  99                 os.Exit(1)\r
 100         } else {\r
 101                 fmt.Println("Created table 'requests' successfully!")\r
 102         }\r
 103 }\r
 104 \r
 105 // an HTTP server that responds to two types of URLs: stats submissions (which it records)\r
 106 // and everything else, which receive a down-page\r
 107 func serve(port int) {\r
 108         requests := 0\r
 109 \r
 110         // routing\r
 111         http.HandleFunc("/", defaultHandler)\r
 112         http.HandleFunc("/stats/submit", makeSubmitHandler(requests))\r
 113         http.Handle("/m/", http.StripPrefix("/m/", http.FileServer(http.Dir("m"))))\r
 114 \r
 115         // serving\r
 116         fmt.Printf("Serving on port %d...\n", port)\r
 117         addr := fmt.Sprintf(":%d", port)\r
 118         http.ListenAndServe(addr, nil)\r
 119 }\r
 120 \r
 121 // intercepts all URLs, displays a landing page\r
 122 func defaultHandler(w http.ResponseWriter, r *http.Request) {\r
 123         err := templates.ExecuteTemplate(w, "landing.html", nil)\r
 124         if err != nil {\r
 125                 http.Error(w, err.Error(), http.StatusInternalServerError)\r
 126         }\r
 127 }\r
 128 \r
 129 // accepts stats requests at a given URL, stores them in requests\r
 130 func makeSubmitHandler(requests int) http.HandlerFunc {\r
 131         return func(w http.ResponseWriter, r *http.Request) {\r
 132                 fmt.Println("in submission handler")\r
 133 \r
 134                 if r.Method != "POST" {\r
 135                         http.Redirect(w, r, "/", http.StatusFound)\r
 136                         return\r
 137                 }\r
 138 \r
 139                 // check for blind ID header. If we don't have it, don't do anything\r
 140                 var blind_id_header string\r
 141                 _, ok := r.Header["X-D0-Blind-Id-Detached-Signature"]\r
 142                 if ok {\r
 143                         fmt.Println("Found a blind_id header. Extracting...")\r
 144                         blind_id_header = r.Header["X-D0-Blind-Id-Detached-Signature"][0]\r
 145                 } else {\r
 146                         fmt.Println("No blind_id header found.")\r
 147                         blind_id_header = ""\r
 148                 }\r
 149 \r
 150                 remoteAddr := getRemoteAddr(r)\r
 151 \r
 152                 // and finally, read the body\r
 153                 body := make([]byte, r.ContentLength)\r
 154                 r.Body.Read(body)\r
 155 \r
 156                 db := getDBConn()\r
 157                 defer db.Close()\r
 158 \r
 159                 _, err := db.Exec("INSERT INTO requests(blind_id_header, ip_addr, body, bodylength) VALUES(?, ?, ?, ?)", blind_id_header, remoteAddr, string(body), r.ContentLength)\r
 160                 if err != nil {\r
 161                         fmt.Println("Unable to insert request.")\r
 162                         fmt.Println(err)\r
 163                 }\r
 164         }\r
 165 }\r
 166 \r
 167 // gets the remote address out of http.Requests with X-Forwarded-For handling\r
 168 func getRemoteAddr(r *http.Request) (remoteAddr string) {\r
 169         val, ok := r.Header["X-Forwarded-For"]\r
 170         if ok {\r
 171                 remoteAddr = val[0]\r
 172         } else {\r
 173                 remoteAddr = r.RemoteAddr\r
 174         }\r
 175 \r
 176         // sometimes a ":<port number>" comes attached, which\r
 177         // needs removing\r
 178         idx := strings.Index(remoteAddr, ":")\r
 179         if idx != -1 {\r
 180                 remoteAddr = remoteAddr[0:idx]\r
 181         }\r
 182 \r
 183         return\r
 184 }\r
 185 \r
 186 // resubmits stats request to a particular URL. this is intended to be used when\r
 187 // you want to write back to the "real" XonStat\r
 188 func resubmit(url string) {\r
 189         db := getDBConn()\r
 190         defer db.Close()\r
 191 \r
 192         rows, err := db.Query("SELECT request_id, ip_addr, blind_id_header, body, bodylength FROM requests ORDER BY request_id")\r
 193         if err != nil {\r
 194                 fmt.Println("Error reading rows from the database. Exiting...")\r
 195                 os.Exit(1)\r
 196         }\r
 197         defer rows.Close()\r
 198 \r
 199         successfulRequests := make([]int, 0, 10)\r
 200         for rows.Next() {\r
 201                 // could use a struct here, but isntead just a bunch of vars\r
 202                 var request_id int\r
 203                 var blind_id_header string\r
 204                 var ip_addr string\r
 205                 var body string\r
 206                 var bodylength int\r
 207 \r
 208                 if err := rows.Scan(&request_id, &ip_addr, &blind_id_header, &body, &bodylength); err != nil {\r
 209                         fmt.Println("Error reading row for submission. Continuing...")\r
 210                         continue\r
 211                 }\r
 212 \r
 213                 req, _ := http.NewRequest("POST", url, strings.NewReader(body))\r
 214                 req.ContentLength = int64(bodylength)\r
 215 \r
 216                 header := map[string][]string{\r
 217                         "X-D0-Blind-Id-Detached-Signature": {blind_id_header},\r
 218                         "X-Forwarded-For":                  {ip_addr},\r
 219                 }\r
 220                 req.Header = header\r
 221 \r
 222                 res, err := http.DefaultClient.Do(req)\r
 223                 if err != nil {\r
 224                         fmt.Printf("Error submitting request #%d. Continuing...\n", request_id)\r
 225                         continue\r
 226                 }\r
 227                 defer res.Body.Close()\r
 228 \r
 229                 fmt.Printf("Request #%d: %s\n", request_id, res.Status)\r
 230 \r
 231                 if res.StatusCode < 500 {\r
 232                         successfulRequests = append(successfulRequests, request_id)\r
 233                 }\r
 234         }\r
 235 \r
 236         // now that we're done resubmitting, let's clean up the successful requests\r
 237         // by deleting them outright from the database\r
 238         for _, val := range successfulRequests {\r
 239                 deleteRequest(db, val)\r
 240         }\r
 241 }\r
 242 \r
 243 // lists all the requests and their information *in the XonStat log format* in\r
 244 // order to 1) show what's in the db and 2) to be able to save/parse it (with\r
 245 // xs_parse) for later use.\r
 246 func list() {\r
 247         db := getDBConn()\r
 248         defer db.Close()\r
 249 \r
 250         rows, err := db.Query("SELECT request_id, ip_addr, blind_id_header, body FROM requests ORDER BY request_id")\r
 251         if err != nil {\r
 252                 fmt.Println("Error reading rows from the database. Exiting...")\r
 253                 os.Exit(1)\r
 254         }\r
 255         defer rows.Close()\r
 256 \r
 257         for rows.Next() {\r
 258                 var request_id int\r
 259                 var blind_id_header string\r
 260                 var ip_addr string\r
 261                 var body string\r
 262 \r
 263                 if err := rows.Scan(&request_id, &ip_addr, &blind_id_header, &body); err != nil {\r
 264                         fmt.Println("Error opening middleman.db. Did you create it?")\r
 265                         continue\r
 266                 }\r
 267 \r
 268                 fmt.Printf("Request: %d\n", request_id)\r
 269                 fmt.Printf("IP Address: %s\n", ip_addr)\r
 270                 fmt.Println("----- BEGIN REQUEST BODY -----")\r
 271 \r
 272                 if len(blind_id_header) > 0 {\r
 273                         fmt.Printf("d0_blind_id: %s\n", blind_id_header)\r
 274                 }\r
 275 \r
 276                 fmt.Print(body)\r
 277                 fmt.Printf("\n----- END REQUEST BODY -----\n")\r
 278         }\r
 279 }\r
 280 \r
 281 // hard-coded sqlite database connection retriever to keep it simple\r
 282 func getDBConn() *sql.DB {\r
 283         conn, err := sql.Open("sqlite3", "./middleman.db")\r
 284 \r
 285         if err != nil {\r
 286                 fmt.Println("Error opening middleman.db. Did you create it?")\r
 287                 os.Exit(1)\r
 288         }\r
 289 \r
 290         return conn\r
 291 }\r
 292 \r
 293 // removes reqeusts from the database by request_id\r
 294 func deleteRequest(db *sql.DB, request_id int) {\r
 295         _, err := db.Exec("delete from requests where request_id = ?", request_id)\r
 296         if err != nil {\r
 297                 fmt.Printf("Could not remove request_id %d from the database. Reason: %v\n", request_id, err)\r
 298         } else {\r
 299                 fmt.Printf("Request #%d removed from the database.\n", request_id)\r
 300         }\r
 301 }\r