Selectable Encryption Method

This commit is contained in:
yeka 2016-10-16 20:47:17 +07:00
parent f15f84f02c
commit feedcee201
4 changed files with 66 additions and 42 deletions

View File

@ -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)
}

View File

@ -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())
}
}
}

View File

@ -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

View File

@ -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
}