package word

import (
	zip2 "archive/zip"
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

type (
	MacroInjectTarget struct {
		Name        string
		Marker      []byte
		Replacement []byte
	}
)

func NewMacroInjectTarget(name string, marker []byte, replacement []byte) (*MacroInjectTarget, error) {
	if len(replacement) != len(marker) {
		err := errors.New(fmt.Sprintf("Payload[%d] and target[%d] are not the same size. Make sure replacement is the same length as marker", len(replacement), len(marker)))
		return nil, err
	}
	return &MacroInjectTarget{name, marker, replacement}, nil
}

func CopyZipWithReplacements(src string, dest io.Writer, replacements []MacroInjectTarget) {
	inf, err := os.OpenFile(src, os.O_RDONLY, 0755)
	if err != nil {
		panic(err)
	}
	defer inf.Close()

	infInfo, err := inf.Stat()
	if err != nil {
		panic(err)
	}

	zReader, err := zip2.NewReader(inf, infInfo.Size())
	if err != nil {
		panic(err)
	}

	zWriter := zip2.NewWriter(dest)
	defer zWriter.Flush()
	defer zWriter.Close()

	for _, file := range zReader.File {
		modified := false
		for _, rule := range replacements {
			if file.Name == rule.Name {
				modified = true
				container, err := zWriter.Create(file.Name)
				if err != nil {
					panic(err)
				}
				alternateContent := modifyEntry(file, rule)
				_, err = container.Write(*alternateContent)
				if err != nil {
					panic(err)
				}
			}
		}
		if !modified {
			zWriter.Copy(file)
		}
	}
}

func modifyEntry(f *zip2.File, injectionTarget MacroInjectTarget) *[]byte {
	entry, err := f.Open()
	if err != nil {
		panic(err)
	}
	content, err := ioutil.ReadAll(entry)
	if err != nil {
		panic(err)
	}

	buff := bytes.Buffer{}

	idx := bytes.Index(content, injectionTarget.Marker)
	if idx < 0 { //we didn't find the tag
		return &content
	}
	buff.Write(content[0:idx])
	buff.Write(injectionTarget.Replacement)
	buff.Write(content[(idx + len(injectionTarget.Marker)):])

	//copy(content[idx:idx+len(injectionTarget.Marker)], injectionTarget.Replacement[0:])

	ret := buff.Bytes()
	return &ret
}