/*
Copyright © 2023 Alvaro Munoz pwntester@github.com

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/
package main

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"encoding/base64"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"time"

	"github.com/GitHubSecurityLab/gh-mrva/cmd"

	"github.com/motemen/go-loghttp"
	"github.com/motemen/go-nuts/roundtime"
)

func main() {
	var transport = &loghttp.Transport{
		Transport:   http.DefaultTransport,
		LogRequest:  LogRequestDump,
		LogResponse: LogResponseDump,
	}

	http.DefaultTransport = transport

	cmd.Execute()
}

func LogRequestDump(req *http.Request) {
	log.Printf(">> %s %s", req.Method, req.URL)
	req.Body = LogBody(req.Body, "request")
}

func LogBody(body io.ReadCloser, from string) io.ReadCloser {

	if body != nil {
		buf, err := io.ReadAll(body)
		if err != nil {
			var w http.ResponseWriter
			log.Fatalf("Error reading %s body: %v", from, err.Error())
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return nil
		}

		IsZipFile := func() bool {
			if len(buf) >= 4 {
				// The header is []byte{ 0x50, 0x4b, 0x03, 0x04 }
				magic := []byte{0x50, 0x4b, 0x03, 0x04}
				if bytes.Equal(buf[0:4], magic) {
					return true
				} else {
					return false
				}
			} else {
				return false
			}
		}

		IsBase64Gzip := func(val []byte) bool {
			// Some important payloads can be listed via
			// base64 -d < foo1 | gunzip | tar t|head -20
			//
			// This function checks the request body up to the `gunzip` part.
			//
			if len(val) >= 4 {
				// Extract header
				hdr := make([]byte, base64.StdEncoding.DecodedLen(4))
				_, err := base64.StdEncoding.Decode(hdr, []byte(val[0:4]))
				if err != nil {
					log.Println("WARNING: IsBase64Gzip decode error:", err)
					return false
				}
				// Check for gzip heading
				magic := []byte{0x1f, 0x8b}
				if bytes.Equal(hdr[0:2], magic) {
					return true
				} else {
					return false
				}
			} else {
				return false
			}
		}

		MaybeJSON := func() bool {
			if len(buf) >= 4 { // {""} is 4 characters
				if bytes.Equal(buf[0:2], []byte("{\"")) {
					return true
				} else {
					return false
				}
			} else {
				return false
			}
		}

		if IsZipFile() {
			// Show index for pk zip archives
			buf1 := make([]byte, len(buf))
			copy(buf1, buf)

			r, err := zip.NewReader(bytes.NewReader(buf1), int64(len(buf1)))
			if err != nil {
				log.Fatal(err)
			}
			// defer r.Close()

			// Print the archive index
			log.Printf(">> %s body:\n", from)
			log.Printf("zip file, contents:\n")

			for _, f := range r.File {
				log.Printf("\t%s\n", f.Name)
			}
		} else if MaybeJSON() {
			// TODO: show index for encoded query packs in the json <value>:
			// {..."query_pack": <value>,...}
			//
			type Message struct {
				// FIXME: exact structure
				ActionRepoRef string   `json:"action_repo_ref"`
				Language      string   `json:"language"`
				QueryPack     string   `json:"query_pack"`
				Repositories  []string `json:"repositories"`
			}
			buf1 := make([]byte, len(buf))
			copy(buf1, buf)
			dec := json.NewDecoder(bytes.NewReader(buf1))
			dec.DisallowUnknownFields()

			var m Message
			if err := dec.Decode(&m); err == io.EOF {
				log.Printf(">> %s body: %v", from, string(buf))
			} else if err != nil {
				log.Printf("WARNING: json decode error: %s\n", err)
				log.Printf(">> %s body: %v", from, string(buf))
			}
			log.Printf(">> %s body:\n", from)
			log.Printf("    \"%s\": \"%s\"\n", "action_repo_ref", m.ActionRepoRef)
			log.Printf("    \"%s\": \"%s\"\n", "language", m.Language)
			log.Printf("    \"%s\": \"%s\"\n", "repositories", m.Repositories[:])

			// Provide custom logging for encoded, compressed tar file
			if IsBase64Gzip([]byte(m.QueryPack)) {
				// These are decoded manually via
				//    base64 -d < foo1 | gunzip | tar t | head -20
				// but we need complete logs for inspection and testing.
				// base64 decode the body
				data, err := base64.StdEncoding.DecodeString(m.QueryPack)
				if err != nil {
					log.Fatalln("body decoding error:", err)
					return nil
				}
				// gunzip the decoded body
				gzb := bytes.NewBuffer(data)
				gzr, err := gzip.NewReader(gzb)
				if err != nil {
					log.Fatal(err)
				}
				// tar t the gunzipped body
				log.Printf("    \"%s\": \n", "query_pack")
				log.Printf("        base64 encoded gzipped tar file, contents:\n")
				tr := tar.NewReader(gzr)
				for {
					hdr, err := tr.Next()
					if err == io.EOF {
						break // End of archive
					}
					if err != nil {
						log.Fatalln("Tar listing failure:", err)
					}
					// TODO: head / tail the listing
					log.Printf("        %s\n", hdr.Name)
				}
			} else {
				log.Printf("    \"%s\": \"%s\"\n", "query_pack", m.QueryPack)
			}
		} else {
			log.Printf(">> %s body: %v", from, string(buf))
		}

		reader := io.NopCloser(bytes.NewBuffer(buf))
		return reader
	}
	return body
}

type contextKey struct {
	name string
}

var ContextKeyRequestStart = &contextKey{"RequestStart"}

func LogResponseDump(resp *http.Response) {
	ctx := resp.Request.Context()
	if start, ok := ctx.Value(ContextKeyRequestStart).(time.Time); ok {
		log.Printf("<< %d %s (%s)", resp.StatusCode, resp.Request.URL,
			roundtime.Duration(time.Since(start), 2))
	} else {
		log.Printf("<< %d %s", resp.StatusCode, resp.Request.URL)
	}

	resp.Body = LogBody(resp.Body, "response")
}