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