szczanieckiej/file/file.go

// SPDX-FileCopyrightText: Adam Evyčędo
//
// SPDX-License-Identifier: AGPL-3.0-or-later

package file

import (
	"archive/tar"
	"archive/zip"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"github.com/ulikunitz/xz"
)

func UnzipGtfs(path string, gtfsFile string) error {
	r, err := zip.OpenReader(filepath.Join(path, gtfsFile))
	if err != nil {
		return err
	}
	defer r.Close()

	for _, f := range r.File {
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer rc.Close()
		txtFile, err := os.Create(filepath.Join(path, f.Name))
		if err != nil {
			return err
		}
		defer txtFile.Close()
		_, err = io.Copy(txtFile, rc)
		if err != nil {
			return err
		}
	}
	return nil
}

func DeleteTxtFiles(path string, gtfsFile string) error {
	txtFiles, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}
	for _, txtFile := range txtFiles {
		if strings.Contains(txtFile.Name(), ".txt") {
			err = os.Remove(filepath.Join(path, txtFile.Name()))
			if err != nil {
				return err
			}
		}
	}
	err = os.Remove(filepath.Join(path, gtfsFile))
	if err != nil {
		return err
	}
	return nil
}

func DeleteBareFiles(path string) error {
	txtFiles, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}
	for _, txtFile := range txtFiles {
		if strings.Contains(txtFile.Name(), ".bare") || strings.Contains(txtFile.Name(), ".lua") {
			err = os.Remove(filepath.Join(path, txtFile.Name()))
			if err != nil {
				log.Printf("while removing bare file %s: %v\n", txtFile.Name(), err)
			}
		}
	}
	return err
}

func CompressBare(path string, gtfsFile string) error {
	trafficFile, err := os.Create(filepath.Join(path, strings.Replace(gtfsFile, ".zip", ".txz", -1)))
	if err != nil {
		return err
	}
	defer trafficFile.Close()
	xzWriter, err := xz.NewWriter(trafficFile)
	if err != nil {
		return err
	}
	defer xzWriter.Close()
	tarWriter := tar.NewWriter(xzWriter)
	defer tarWriter.Close()

	files := []string{"calendar.bare", "ix_stop_codes.bare", "ix_stop_names.bare", "lines.bare", "stops.bare", "trips.bare", "vehicles.bare", "ix_lines.bare", "ix_line_codes.bare", "feed_info.bare", "agencies.bare", "ix_trips.bare", "updates.lua", "vehicles.lua", "alerts.lua"}
	optionalFiles := map[string]struct{}{
		"updates.lua":  {},
		"vehicles.lua": {},
		"alerts.lua":   {},
	}
	for _, file := range files {
		stat, err := os.Stat(filepath.Join(path, file))
		if err != nil {
			if _, ok := optionalFiles[file]; ok {
				log.Printf("while stating %s: %v; ignoring\n", file, err)
				continue
			}
			return err
		}
		bareFile, err := os.Open(filepath.Join(path, file))
		if err != nil {
			return err
		}
		defer bareFile.Close()
		hdr := &tar.Header{
			Name: file,
			Mode: 0600,
			Size: stat.Size(),
		}
		if err := tarWriter.WriteHeader(hdr); err != nil {
			return err
		}
		if _, err := io.Copy(tarWriter, bareFile); err != nil {
			return err
		}
	}
	return nil
}

func MoveTraffic(srcName, dstName, path, feedHome string) error {
	inputFile, err := os.Open(filepath.Join(path, strings.Replace(srcName, ".zip", ".txz", -1)))
	if err != nil {
		log.Printf("Couldn't open source file: %s\n", err)
		return err
	}
	outputFile, err := os.Create(filepath.Join(feedHome, dstName))
	if err != nil {
		inputFile.Close()
		log.Printf("Couldn't open dest file: %s\n", err)
		return err
	}
	defer outputFile.Close()
	_, err = io.Copy(outputFile, inputFile)
	inputFile.Close()
	if err != nil {
		log.Printf("Writing to output file failed: %s\n", err)
		return err
	}
	err = os.Remove(filepath.Join(path, strings.Replace(srcName, ".zip", ".txz", -1)))
	if err != nil {
		log.Printf("Failed removing original file: %s\n", err)
		return err
	}
	return nil
}

func ListVersions(path string) ([]string, error) {
	versions := []string{}
	trafficFiles, err := ioutil.ReadDir(path)
	if err != nil {
		return nil, err
	}
	for _, txtFile := range trafficFiles {
		if strings.Contains(txtFile.Name(), ".txz") {
			versionName := strings.Replace(txtFile.Name(), ".txz", "", 1)
			versions = append(versions, versionName)
		}
	}
	return versions, nil
}

func UnpackTraffic(dataHome, feedName string) error {
	path := filepath.Join(dataHome, feedName)
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}
	unpacked := map[string]bool{}
	packed := map[string]bool{}
	for _, file := range files {
		if strings.Contains(file.Name(), ".txz") {
			packed[strings.Replace(file.Name(), ".txz", "", 1)] = true
		} else {
			unpacked[file.Name()] = true
		}
	}

	for version := range packed {
		if !unpacked[version] {
			trafficPath := filepath.Join(path, version) + ".txz"
			versionPath := filepath.Join(path, version)
			err := os.Mkdir(versionPath, 0755)
			if err != nil {
				return err
			}
			trafficFile, err := os.Open(trafficPath)
			if err != nil {
				return err
			}
			defer trafficFile.Close()
			xzReader, err := xz.NewReader(trafficFile)
			if err != nil {
				return err
			}
			tarReader := tar.NewReader(xzReader)
			for {
				hdr, err := tarReader.Next()
				if err == io.EOF {
					break
				}
				if err != nil {
					return err
				}
				barePath := filepath.Join(versionPath, hdr.Name)
				bareFile, err := os.Create(barePath)
				if err != nil {
					return err
				}
				defer bareFile.Close()
				if _, err := io.Copy(bareFile, tarReader); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func CleanOldVersions(path string, validVersions map[string]bool) error {
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	for _, file := range files {
		name := file.Name()
		versionString := strings.Replace(name, ".txz", "", 1)
		if !validVersions[versionString] {
			filePath := filepath.Join(path, name)
			os.RemoveAll(filePath)
		}
	}
	return nil
}