reforger_utils/util.go

package reforger_utils

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
)

var Debug = false
var DownThreads = 8

func DoAssetsRequest(addonID, version string) (AssetsReply, error) {
	url := "https://api-ar-workshop.bistudio.com/workshop-api/api/v3.0/s2s/assets/download-list"
	body := `{"assets":{"` + addonID + `":"` + version + `"}}`
	req, err := http.NewRequest("POST", url, strings.NewReader(body))
	if err != nil {
		return AssetsReply{}, err
	}

	req.Header.Add("x-client-id", "$edb1b7862bba5cade1f6e06bfdeac2c")
	req.Header.Add("x-client-secret", "$8b415ea2aa11bd51f2f5b5a9dcb8476")
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("user-agent", "Arma Reforger/1.4.0.53 (Headless; Windows)")
	req.Header.Add("content-length", fmt.Sprintf("%d", len(body)))

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return AssetsReply{}, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		// print entire body
		body, _ := io.ReadAll(res.Body)
		fmt.Println(string(body))
		return AssetsReply{}, fmt.Errorf("bad status: %s", res.Status)
	}

	// Process response
	decoder := json.NewDecoder(res.Body)
	var assetsReply AssetsReply
	if err := decoder.Decode(&assetsReply); err != nil {
		panic(err)
	}

	return assetsReply, nil
}

func GetManifest(manifestUrl string) (Manifest, error) {
	url := "https://ar-gcp-cdn.bistudio.com/manifest/" + manifestUrl
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return Manifest{}, err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return Manifest{}, err
	}
	defer res.Body.Close()

	if res.StatusCode != 200 {
		return Manifest{}, fmt.Errorf("bad status: %s", res.Status)
	}

	// Decode the JSON response
	var manifest Manifest
	if err := json.NewDecoder(res.Body).Decode(&manifest); err != nil {
		panic(err)
	}

	return manifest, nil
}

func Download(manifest Manifest, downloadDir string, filen string) ([]byte, error) {

	if Debug {
		// create the directory if it doesn't exist
		if err := os.MkdirAll(downloadDir+"/fragments/"+filen, os.ModePerm); err != nil {
			return nil, err
		}
	}

	// download remainder
	downloaded := 0
	remContent := make([]byte, 0)
	currentRem := 0
	if manifest.Remainder.Size > 0 {
		url := transformShaToURL(manifest.Remainder.Sha512, manifest.Remainder.Size)
		content, err := getContent(url)
		if err != nil {
			return nil, err
		}
		remContent = content
		downloaded += len(content)

		if Debug {
			// save remainder to file in the fragments/ directory, prefix with gap
			if err := os.WriteFile(filepath.Join(downloadDir, "fragments", filen, fmt.Sprintf("%d.%s.%d.bytes", manifest.Remainder.Offsets[0], manifest.Remainder.Sha512, manifest.Remainder.Size)), remContent, os.ModePerm); err != nil {
				return nil, err
			}
		}
	}

	trackFrag := make([]Frag, 0)

	// start downloading chunks
	ret := make([]byte, manifest.Size)
	totalMB := float32(manifest.Size) / 1024.0 / 1024.0

	var wg sync.WaitGroup
	var mu sync.Mutex

	fragmentsChan := make(chan Fragment, len(manifest.Fragments))
	errChan := make(chan error, DownThreads)

	// Start worker goroutines
	for i := 0; i < DownThreads; i++ {
		wg.Go(func() {
			for fragment := range fragmentsChan {
				url := transformShaToURL(fragment.Sha512, fragment.Size)
				content, err := getContent(url)
				if err != nil {
					// Handle error, maybe by sending it to an error channel
					// For now, we'll just skip and print.
					fmt.Printf("Error downloading fragment %s: %v\n", fragment.Sha512, err)
					errChan <- err
				}

				mu.Lock()
				// Safely write to the shared slice and update progress
				for _, offset := range fragment.Offsets {
					copy(ret[offset:offset+len(content)], content)
				}
				downloaded += len(content) * len(fragment.Offsets)

				// Track downloaded fragments
				for _, offset := range fragment.Offsets {
					trackFrag = append(trackFrag, Frag{
						Offset: offset,
						Size:   len(content),
					})
				}

				// Print progress
				downedMB := float32(downloaded) / 1024.0 / 1024.0
				percent := float32(downloaded) / float32(manifest.Size) * 100.0
				fmt.Printf("\r\033[32mDownloaded\033[0m %.2fMB/%.2fMB (\033[36m%.2f%%\033[0m) of \033[33m%s\033[0m", downedMB, totalMB, percent, filen)

				mu.Unlock()

				// The Debug section can be safely moved here as well
				if os.Getenv("DEBUG") == "true" { // Using env var for a cleaner check
					mu.Lock() // Re-lock for file write
					if err := os.WriteFile(filepath.Join(downloadDir, "fragments", filen, fmt.Sprintf("%d.%s.%d.bytes", fragment.Offsets[0], fragment.Sha512, fragment.Size)), content, os.ModePerm); err != nil {
						fmt.Printf("Error saving file: %v\n", err)
					}
					mu.Unlock()
				}
			}
		})
	}

	// Send fragments to the channel
	for _, fragment := range manifest.Fragments {
		fragmentsChan <- fragment
	}
	close(fragmentsChan) // Close the channel to signal workers no more tasks are coming

	// Wait for all workers to finish
	wg.Wait()

	// Check for errors
	if len(errChan) > 0 {
		return nil, <-errChan // Return the first error encountered
	}
	close(errChan)

	// sort trackFrag by Offset
	slices.SortFunc(trackFrag, func(a, b Frag) int {
		return a.Offset - b.Offset
	})

	// check for gaps in trackFrag and fill with remainder content
	curPos := 0

	for _, frag := range trackFrag {
		if frag.Offset > curPos {
			// gap detected
			gapSize := frag.Offset - curPos
			if currentRem+gapSize > len(remContent) {
				panic("not enough remainder content to fill gap")
			}
			copy(ret[curPos:curPos+gapSize], remContent[currentRem:currentRem+gapSize])
			remContent = remContent[gapSize:]
		}
		curPos = frag.Offset + frag.Size
	}

	// Append rest of remainder content
	copy(ret[curPos:], remContent)
	fmt.Println()

	return ret, nil
}

func getContent(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("bad status: %s", resp.Status)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return body, nil
}

func transformShaToURL(sha512s string, size int) string {
	return "https://ar-gcp-cdn.bistudio.com/fragment/" + sha512s[0:4] + "/" + sha512s[4:8] + "/" + sha512s[8:] + "/" + fmt.Sprintf("%d", size)
}