diff --git a/crypto.go b/crypto.go index 137b674..3f86df1 100644 --- a/crypto.go +++ b/crypto.go @@ -19,7 +19,14 @@ import ( "golang.org/x/crypto/pbkdf2" ) +type EncryptionMethod int + const ( + ZipStandardEncryption EncryptionMethod = 1 + AES128Encryption EncryptionMethod = 2 + AES192Encryption EncryptionMethod = 3 + AES256Encryption EncryptionMethod = 4 + // AES key lengths aes128 = 16 aes192 = 24 @@ -379,13 +386,14 @@ func encryptStream(key []byte, w io.Writer) (io.Writer, error) { // newEncryptionWriter returns an io.Writer that when written to, 1. writes // out the salt, 2. writes out pwv, and 3. writes out authenticated, encrypted // data. The authcode will be written out in fileWriter.close(). -func newEncryptionWriter(w io.Writer, password passwordFn, fw *fileWriter) (io.Writer, error) { - var salt [16]byte +func newEncryptionWriter(w io.Writer, password passwordFn, fw *fileWriter, aesstrength byte) (io.Writer, error) { + keysize := aesKeyLen(aesstrength) + salt := make([]byte, keysize / 2) _, err := rand.Read(salt[:]) if err != nil { return nil, errors.New("zip: unable to generate random salt") } - ekey, akey, pwv := generateKeys(password(), salt[:], aes256) + ekey, akey, pwv := generateKeys(password(), salt[:], keysize) fw.hmac = hmac.New(sha1.New, akey) aw := &authWriter{ hmac: fw.hmac, @@ -424,11 +432,23 @@ func (h *FileHeader) writeWinZipExtra() { eb.uint16(7) // following data size is 7 eb.uint16(2) // ae 2 eb.uint16(0x4541) // "AE" - eb.uint8(3) // aes256 + eb.uint8(h.aesStrength) // aes256 eb.uint16(h.Method) // original compression method h.Extra = append(h.Extra, buf[:]...) } +func (h *FileHeader) setEncryptionMethod(enc EncryptionMethod) { + h.encryption = enc + switch enc { + case AES128Encryption: + h.aesStrength = 1 + case AES192Encryption: + h.aesStrength = 2 + case AES256Encryption: + h.aesStrength = 3 + } +} + func (h *FileHeader) setEncryptionBit() { h.Flags |= 0x1 } @@ -452,11 +472,12 @@ type passwordFn func() []byte // contents will be encrypted with AES-256 using the given password. The // file's contents must be written to the io.Writer before the next call // to Create, CreateHeader, or Close. -func (w *Writer) Encrypt(name string, password string) (io.Writer, error) { +func (w *Writer) Encrypt(name string, password string, enc EncryptionMethod) (io.Writer, error) { fh := &FileHeader{ Name: name, Method: Deflate, } fh.SetPassword(password) + fh.setEncryptionMethod(enc) return w.CreateHeader(fh) } diff --git a/crypto_test.go b/crypto_test.go index e7b7277..0c75234 100644 --- a/crypto_test.go +++ b/crypto_test.go @@ -175,42 +175,44 @@ func TestPasswordWriteSimple(t *testing.T) { contents := []byte("Hello World") conLen := len(contents) - raw := new(bytes.Buffer) - zipw := NewWriter(raw) - w, err := zipw.Encrypt("hello.txt", "golang") - if err != nil { - t.Errorf("Expected to create a new FileHeader") - } - n, err := io.Copy(w, bytes.NewReader(contents)) - if err != nil || n != int64(conLen) { - t.Errorf("Expected to write the full contents to the writer.") - } - zipw.Close() + for _, enc := range []EncryptionMethod{AES128Encryption, AES192Encryption, AES256Encryption} { + raw := new(bytes.Buffer) + zipw := NewWriter(raw) + w, err := zipw.Encrypt("hello.txt", "golang", enc) + if err != nil { + t.Errorf("Expected to create a new FileHeader") + } + n, err := io.Copy(w, bytes.NewReader(contents)) + if err != nil || n != int64(conLen) { + t.Errorf("Expected to write the full contents to the writer.") + } + zipw.Close() - // Read the zip - buf := new(bytes.Buffer) - zipr, err := NewReader(bytes.NewReader(raw.Bytes()), int64(raw.Len())) - if err != nil { - t.Errorf("Expected to open a new zip reader: %v", err) - } - nn := len(zipr.File) - if nn != 1 { - t.Errorf("Expected to have one file in the zip archive, but has %d files", nn) - } - z := zipr.File[0] - z.SetPassword("golang") - rr, err := z.Open() - if err != nil { - t.Errorf("Expected to open the readcloser: %v", err) - } - n, err = io.Copy(buf, rr) - if err != nil { - t.Errorf("Expected to write to temporary buffer: %v", err) - } - if n != int64(conLen) { - t.Errorf("Expected to copy %d bytes to temp buffer, but copied %d bytes instead", conLen, n) - } - if !bytes.Equal(contents, buf.Bytes()) { - t.Errorf("Expected the unzipped contents to equal '%s', but was '%s' instead", contents, buf.Bytes()) + // Read the zip + buf := new(bytes.Buffer) + zipr, err := NewReader(bytes.NewReader(raw.Bytes()), int64(raw.Len())) + if err != nil { + t.Errorf("Expected to open a new zip reader: %v", err) + } + nn := len(zipr.File) + if nn != 1 { + t.Errorf("Expected to have one file in the zip archive, but has %d files", nn) + } + z := zipr.File[0] + z.SetPassword("golang") + rr, err := z.Open() + if err != nil { + t.Errorf("Expected to open the readcloser: %v", err) + } + n, err = io.Copy(buf, rr) + if err != nil { + t.Errorf("Expected to write to temporary buffer: %v", err) + } + if n != int64(conLen) { + t.Errorf("Expected to copy %d bytes to temp buffer, but copied %d bytes instead", conLen, n) + } + if !bytes.Equal(contents, buf.Bytes()) { + t.Errorf("Expected the unzipped contents to equal '%s', but was '%s' instead", contents, buf.Bytes()) + } } } diff --git a/struct.go b/struct.go index 1e0cba2..7e50a5c 100644 --- a/struct.go +++ b/struct.go @@ -101,6 +101,7 @@ type FileHeader struct { // https://www.imperialviolet.org/2015/05/16/aeads.html DeferAuth bool + encryption EncryptionMethod password passwordFn // Returns the password to use when reading/writing ae uint16 aesStrength byte diff --git a/writer.go b/writer.go index 27d125e..cc08a3e 100644 --- a/writer.go +++ b/writer.go @@ -232,7 +232,7 @@ func (w *Writer) CreateHeader(fh *FileHeader) (io.Writer, error) { // we have a password and need to encrypt. fh.writeWinZipExtra() fh.Method = 99 // ok to change, we've gotten the comp and wrote extra - ew, err := newEncryptionWriter(sw, fh.password, fw) + ew, err := newEncryptionWriter(sw, fh.password, fw, fh.aesStrength) if err != nil { return nil, err }