github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/uzip/uzip.go (about)

     1  // Copyright 2018 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package uzip contains utilities for file system->zip and zip->file system conversions.
     6  package uzip
     7  
     8  import (
     9  	"archive/zip"
    10  	"fmt"
    11  	"io"
    12  	"os"
    13  	"path/filepath"
    14  	"strings"
    15  )
    16  
    17  // ToZip packs the all files at dir to a zip archive at dest.
    18  func ToZip(dir, dest string) error {
    19  	info, err := os.Stat(dir)
    20  	if err != nil {
    21  		return err
    22  	}
    23  	if !(info.IsDir()) {
    24  		return fmt.Errorf("%s is not a directory", dir)
    25  	}
    26  	archive, err := os.Create(dest)
    27  	if err != nil {
    28  		return err
    29  	}
    30  	defer archive.Close()
    31  
    32  	z := zip.NewWriter(archive)
    33  	defer z.Close()
    34  
    35  	err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
    36  		if err != nil {
    37  			return err
    38  		}
    39  
    40  		// do not include srcDir into archive
    41  		if info.Name() == filepath.Base(dir) {
    42  			return nil
    43  		}
    44  
    45  		header, err := zip.FileInfoHeader(info)
    46  		if err != nil {
    47  			return err
    48  		}
    49  
    50  		// adjust header.Name to preserve folder strulture
    51  		header.Name = strings.TrimPrefix(path, dir)
    52  
    53  		if info.IsDir() {
    54  			header.Name += "/"
    55  		} else {
    56  			header.Method = zip.Deflate
    57  		}
    58  
    59  		writer, err := z.CreateHeader(header)
    60  		if err != nil {
    61  			return err
    62  		}
    63  
    64  		if info.IsDir() {
    65  			return nil
    66  		}
    67  
    68  		file, err := os.Open(path)
    69  		if err != nil {
    70  			return err
    71  		}
    72  		defer file.Close()
    73  		_, err = io.Copy(writer, file)
    74  		return err
    75  	})
    76  
    77  	return err
    78  }
    79  
    80  // FromZip extracts the zip archive at src to dir.
    81  func FromZip(src, dir string) error {
    82  	z, err := zip.OpenReader(src)
    83  	if err != nil {
    84  		return err
    85  	}
    86  
    87  	if err = os.MkdirAll(dir, 0755); err != nil {
    88  		return err
    89  	}
    90  
    91  	for _, file := range z.File {
    92  		path := filepath.Join(dir, file.Name)
    93  		if file.FileInfo().IsDir() {
    94  			if err = os.MkdirAll(path, file.Mode()); err != nil {
    95  				return err
    96  			}
    97  			continue
    98  		}
    99  
   100  		fileReader, err := file.Open()
   101  		if err != nil {
   102  			return err
   103  		}
   104  
   105  		targetFile, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, file.Mode())
   106  		if err != nil {
   107  			return err
   108  		}
   109  
   110  		if _, err := io.Copy(targetFile, fileReader); err != nil {
   111  			return err
   112  		}
   113  
   114  		if err = fileReader.Close(); err != nil {
   115  			return err
   116  		}
   117  
   118  		if err = targetFile.Close(); err != nil {
   119  			return err
   120  		}
   121  	}
   122  
   123  	return nil
   124  }