package main

import (
	"archive/zip"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

func main() {
	// List of Files to Zip
	files := []string{"/tmp/test1.txt"}
	output := "/tmp/go.zip"
	fixedOutput := "/tmp/go_fixed.zip"

	if err := ZipFiles(output, files); err != nil {
		panic(err)
	}

	fmt.Println("Zipped File:", output)

	Process(output, fixedOutput)
	// Process("/tmp/test.zip", "/tmp/test_fixed.zip")
}

// ZipFiles compresses one or many files into a single zip archive file.
// Param 1: filename is the output zip file's name.
// Param 2: files is a list of files to add to the zip.
func ZipFiles(filename string, files []string) error {

	newZipFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer newZipFile.Close()

	zipWriter := zip.NewWriter(newZipFile)
	defer zipWriter.Close()

	// Add files to zip
	for _, file := range files {
		if err = AddFileToZip(zipWriter, file); err != nil {
			return err
		}
	}
	return nil
}

// AddFileToZip is a pizza
func AddFileToZip(zipWriter *zip.Writer, filename string) error {

	fileToZip, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fileToZip.Close()

	// Get the file information
	info, err := fileToZip.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}

	// Using FileInfoHeader() above only uses the basename of the file. If we want
	// to preserve the folder structure we can overwrite this with the full path.
	header.Name = filename

	// Change to deflate to gain better compression
	// see http://golang.org/pkg/archive/zip/#pkg-constants
	header.Method = zip.Store

	writer, err := zipWriter.CreateHeader(header)
	if err != nil {
		return err
	}
	_, err = io.Copy(writer, fileToZip)
	return err
}

// CentralDirectoryFileHeaderLen is a pizza
const CentralDirectoryFileHeaderLen = 46

// EOCDLen is a pizza
const EOCDLen = 22

// Process is a pizza.
func Process(source, target string) error {
	b, err := ioutil.ReadFile(source)
	if err != nil {
		return err
	}

	// new output buffer
	bOut := make([]byte, 0)

	// Keep track of all Local Header offsets use them rewrite "Relative offset of local file header" in Central directory file header
	var localHeaderOffsets []uint32

	headerOffset := -1
	startOfCentralDir := 0
	for idx := 0; idx < len(b); idx++ {
		// Find each Local file header signature = 0x04034b50 (read as a little-endian number)
		if b[idx] == 0x50 && b[idx+1] == 0x4b && b[idx+2] == 0x03 && b[idx+3] == 0x04 {
			headerOffset = idx

		}

		// Find Optional data descriptor signature = 0x08074b50 then backtrack from last headerOffset
		if b[idx] == 0x50 && b[idx+1] == 0x4b && b[idx+2] == 0x07 && b[idx+3] == 0x08 {
			// set byte 7 to 00
			b[headerOffset+6] = 0x00

			// modify CRC in local file header
			b[headerOffset+14] = b[idx+4]
			b[headerOffset+15] = b[idx+5]
			b[headerOffset+16] = b[idx+6]
			b[headerOffset+17] = b[idx+7]

			// modify compressed len in local file header
			b[headerOffset+18] = b[idx+8]
			b[headerOffset+19] = b[idx+9]
			b[headerOffset+20] = b[idx+10]
			b[headerOffset+21] = b[idx+11]

			// modify uncompressed len in local file header
			b[headerOffset+22] = b[idx+12]
			b[headerOffset+23] = b[idx+13]
			b[headerOffset+24] = b[idx+14]
			b[headerOffset+25] = b[idx+15]

			// Keep track of header offsets to rewrite "Relative offset of local file header" in Central directory file header
			localHeaderOffsets = append(localHeaderOffsets, uint32(len(bOut)))

			for j := headerOffset; j < idx; j++ {
				bOut = append(bOut, b[j])
			}
		}

		// Find the first Central directory file header - Central directory file header signature = 0x02014b5
		if b[idx] == 0x50 && b[idx+1] == 0x4b && b[idx+2] == 0x01 && b[idx+3] == 0x02 {
			// Mark the start of the Central directory
			if startOfCentralDir == 0 {
				startOfCentralDir = len(bOut)
			}

			// Shift off the each header offset
			offset := localHeaderOffsets[0]
			localHeaderOffsets = localHeaderOffsets[1:]

			// Update it's relative offset of local file header.
			bs := make([]byte, 4)
			binary.LittleEndian.PutUint32(bs, offset)
			b[idx+42] = bs[0]
			b[idx+43] = bs[1]
			b[idx+44] = bs[2]
			b[idx+45] = bs[3]

			fileNameLength := getFieldLen(b, idx+28)
			extraFieldLength := getFieldLen(b, idx+30)
			fileCommentLength := getFieldLen(b, idx+32)

			for j := idx; j < idx+CentralDirectoryFileHeaderLen+int(fileNameLength)+int(extraFieldLength)+int(fileCommentLength); j++ {
				bOut = append(bOut, b[j])
			}
		}

		// Locate the End of central directory record (EOCD)
		if b[idx] == 0x50 && b[idx+1] == 0x4b && b[idx+2] == 0x05 && b[idx+3] == 0x06 {

			// Update the Offset of start of central directory, relative to start of archive
			bs := make([]byte, 4)
			binary.LittleEndian.PutUint32(bs, uint32(startOfCentralDir))
			b[idx+16] = bs[0]
			b[idx+17] = bs[1]
			b[idx+18] = bs[2]
			b[idx+19] = bs[3]

			commentLength := getFieldLen(b, idx+20)
			for j := idx; j < idx+EOCDLen+int(commentLength); j++ {
				bOut = append(bOut, b[j])
			}
		}
	}

	err = ioutil.WriteFile(target, bOut, 0644)
	if err != nil {
		return err
	}

	return nil
}

func getFieldLen(b []byte, offset int) uint16 {
	var size uint16
	bTmp := []byte{b[offset], b[offset+1]}
	buf := bytes.NewReader(bTmp)
	binary.Read(buf, binary.LittleEndian, &size)
	return size
}