//
// this code use the idea from this https://gist.github.com/hvoecking/10772475
// list of xss payload https://github.com/payloadbox/xss-payload-list

package main

import (
	"fmt"
	"log"
	"reflect"

	"github.com/microcosm-cc/bluemonday"
)

// traverse through struct
// refer from this cool snippet https://gist.github.com/hvoecking/10772475
func sanitize(p *bluemonday.Policy, obj interface{}) interface{} {
	// Wrap the original in a reflect.Value
	original := reflect.ValueOf(obj)

	copy := reflect.New(original.Type()).Elem()
	sanitizeRecursive(p, copy, original)

	// Remove the reflection wrapper
	return copy.Interface()
}

func sanitizeRecursive(p *bluemonday.Policy, copy, original reflect.Value) {
	switch original.Kind() {
	// The first cases handle nested structures and sanitize them recursively

	// If it is a pointer we need to unwrap and call once again
	case reflect.Ptr:
		// To get the actual value of the original we have to call Elem()
		// At the same time this unwraps the pointer so we don't end up in
		// an infinite recursion
		originalValue := original.Elem()
		// Check if the pointer is nil
		if !originalValue.IsValid() {
			return
		}
		// Allocate a new object and set the pointer to it
		copy.Set(reflect.New(originalValue.Type()))
		// Unwrap the newly created pointer
		sanitizeRecursive(p, copy.Elem(), originalValue)

	// If it is an interface (which is very similar to a pointer), do basically the
	// same as for the pointer. Though a pointer is not the same as an interface so
	// note that we have to call Elem() after creating a new object because otherwise
	// we would end up with an actual pointer
	case reflect.Interface:
		// Get rid of the wrapping interface
		originalValue := original.Elem()
		// Create a new object. Now new gives us a pointer, but we want the value it
		// points to, so we have to call Elem() to unwrap it
		copyValue := reflect.New(originalValue.Type()).Elem()
		sanitizeRecursive(p, copyValue, originalValue)
		copy.Set(copyValue)

	// If it is a struct we sanitize each field
	case reflect.Struct:
		for i := 0; i < original.NumField(); i += 1 {
			sanitizeRecursive(p, copy.Field(i), original.Field(i))
		}

	// If it is a slice we create a new slice and sanitize each element
	case reflect.Slice:
		copy.Set(reflect.MakeSlice(original.Type(), original.Len(), original.Cap()))
		for i := 0; i < original.Len(); i += 1 {
			sanitizeRecursive(p, copy.Index(i), original.Index(i))
		}

	// If it is a map we create a new map and sanitize each value
	case reflect.Map:
		copy.Set(reflect.MakeMap(original.Type()))
		for _, key := range original.MapKeys() {
			originalValue := original.MapIndex(key)
			// New gives us a pointer, but again we want the value
			copyValue := reflect.New(originalValue.Type()).Elem()
			sanitizeRecursive(p, copyValue, originalValue)
			copy.SetMapIndex(key, copyValue)
		}

	// Otherwise we cannot traverse anywhere so this finishes the the recursion

	// If it is a string sanitize it (yay finally we're doing what we came for)
	case reflect.String:
		sanitizedString := p.Sanitize(original.Interface().(string))
		copy.SetString(sanitizedString)

	// And everything else will simply be taken from the original
	default:
		copy.Set(original)
	}

}

func main() {
	log.SetFlags(log.Lshortfile | log.LstdFlags)

	// Do this once for each unique policy, and use the policy for the life of the program
	// Policy creation/editing is not safe to use in multiple goroutines
	p := bluemonday.UGCPolicy()

	type Bazz struct {
		D string
		E map[interface{}]interface{}
	}

	type Bar struct {
		B        string
		C        map[interface{}]interface{}
		FBazz    Bazz
		FBazzPtr *Bazz
	}

	type Foo struct {
		Fstr1   string
		Fstr2   string
		Fint    int
		Fbar    Bar
		FbarPtr *Bar
	}

	foo := Foo{
		Fstr1: `<a onblur="alert(secret)" href="http://www.google.com">Google</a>`,
		Fstr2: "safe string",
		Fint:  100,
		Fbar: Bar{
			B: `<a onblur="alert(secret)" href="http://struct.com">xyz</a>`,
			C: map[interface{}]interface{}{"cField1": 10, "cField2": `<a onblur="alert(secret)" href="http://www.google.com">Google</a>`},
			FBazz: Bazz{
				D: `abc`,
				E: map[interface{}]interface{}{
					1: `abc`,
				},
			},
			FBazzPtr: &Bazz{
				D: `abc`,
				E: map[interface{}]interface{}{
					1: `<image/src/onerror=prompt(8)>`,
				},
			},
		},
		FbarPtr: &Bar{
			B: `<a onblur="alert(secret)" href="http://pointer.com">xyz</a>`,
			C: map[interface{}]interface{}{"cField1": 10, "cField2": `<a onblur="alert(secret)" href="http://www.google.com">Google</a>`},
			FBazz: Bazz{
				D: `abc`,
				E: map[interface{}]interface{}{
					1: `<image/src/onerror=prompt(8)>`,
				},
			},
			FBazzPtr: &Bazz{
				D: `abc`,
				E: map[interface{}]interface{}{
					1: `<a onblur="alert(secret)" href="http://pointer.com">xyz</a>`,
				},
			},
		},
	}

	fmt.Println("<<<<<<<<<<<<<<<<<<")

	log.Println("before %+v", foo)

	fmt.Println(">>>>>>>>>>>>>>>>>>")

	fooo := sanitize(p, foo)

	log.Println("after sanitized: %+v", fooo)
	log.Println("after sanitized: %+v", fooo.(Foo).Fbar)
	log.Println("after sanitized: %+v", fooo.(Foo).Fbar.FBazzPtr)
	log.Println("after sanitized: %+v", fooo.(Foo).FbarPtr)
	log.Println("after sanitized: %+v", fooo.(Foo).FbarPtr.FBazzPtr)
	fmt.Println("==================")
}