1package steg
2
3import (
4	"bytes"
5	"encoding/binary"
6	"fmt"
7	"github.com/DimitarPetrov/stegify/bits"
8	"image"
9	"image/draw"
10	_ "image/jpeg" //register jpeg image format
11	"image/png"
12	"io"
13	"io/ioutil"
14	"os"
15)
16
17const dataSizeHeaderReservedBytes = 20 // 20 bytes results in 30 usable bits
18
19//Encode performs steganography encoding of data Reader in carrier
20//and writes it to the result Writer encoded as PNG image.
21func Encode(carrier io.Reader, data io.Reader, result io.Writer) error {
22	RGBAImage, format, err := getImageAsRGBA(carrier)
23	if err != nil {
24		return fmt.Errorf("error parsing carrier image: %v", err)
25	}
26
27	dataBytes := make(chan byte, 128)
28	errChan := make(chan error)
29
30	go readData(data, dataBytes, errChan)
31
32	dx := RGBAImage.Bounds().Dx()
33	dy := RGBAImage.Bounds().Dy()
34
35	hasMoreBytes := true
36
37	var count int
38	var dataCount uint32
39
40	for x := 0; x < dx && hasMoreBytes; x++ {
41		for y := 0; y < dy && hasMoreBytes; y++ {
42			if count >= dataSizeHeaderReservedBytes {
43				c := RGBAImage.RGBAAt(x, y)
44				hasMoreBytes, err = setColorSegment(&c.R, dataBytes, errChan)
45				if err != nil {
46					return err
47				}
48				if hasMoreBytes {
49					dataCount++
50				}
51				hasMoreBytes, err = setColorSegment(&c.G, dataBytes, errChan)
52				if err != nil {
53					return err
54				}
55				if hasMoreBytes {
56					dataCount++
57				}
58				hasMoreBytes, err = setColorSegment(&c.B, dataBytes, errChan)
59				if err != nil {
60					return err
61				}
62				if hasMoreBytes {
63					dataCount++
64				}
65				RGBAImage.SetRGBA(x, y, c)
66			} else {
67				count += 4
68			}
69		}
70	}
71
72	select {
73	case _, ok := <-dataBytes: // if there is more data
74		if ok {
75			return fmt.Errorf("data file too large for this carrier")
76		}
77	default:
78	}
79
80	setDataSizeHeader(RGBAImage, quartersOfBytesOf(dataCount))
81
82	switch format {
83	case "png", "jpeg":
84		return png.Encode(result, RGBAImage)
85	default:
86		return fmt.Errorf("unsupported carrier format")
87	}
88}
89
90//MultiCarrierEncode performs steganography encoding of data Reader in equal pieces in each of the carriers
91//and writes it to the result Writers encoded as PNG images.
92func MultiCarrierEncode(carriers []io.Reader, data io.Reader, results []io.Writer) error {
93	if len(carriers) != len(results) {
94		return fmt.Errorf("different number of carriers and results")
95	}
96
97	dataBytes, err := ioutil.ReadAll(data)
98	if err != nil {
99		return fmt.Errorf("error reading data %v", err)
100	}
101
102	chunkSize := len(dataBytes) / len(carriers)
103	dataChunks := make([]io.Reader, 0, len(carriers))
104	chunksCount := 0
105	for i := 0; i < len(dataBytes) && chunksCount < len(carriers); i += chunkSize {
106		chunksCount++
107		if i+chunkSize >= len(dataBytes) || chunksCount == len(carriers) { // last iteration
108			dataChunks = append(dataChunks, bytes.NewReader(dataBytes[i:]))
109		}
110		dataChunks = append(dataChunks, bytes.NewReader(dataBytes[i:i+chunkSize]))
111	}
112
113	for i := 0; i < len(carriers); i++ {
114		if err := Encode(carriers[i], dataChunks[i], results[i]); err != nil {
115			return fmt.Errorf("error encoding chunk with index %d: %v", i, err)
116		}
117	}
118	return nil
119}
120
121//EncodeByFileNames performs steganography encoding of data file in carrier file
122//and saves the steganography encoded product in new file.
123func EncodeByFileNames(carrierFileName, dataFileName, resultFileName string) (err error) {
124	return MultiCarrierEncodeByFileNames([]string{carrierFileName}, dataFileName, []string{resultFileName})
125}
126
127//MultiCarrierEncodeByFileNames performs steganography encoding of data file in equal pieces in each of the carrier files
128//and saves the steganography encoded product in new set of result files.
129func MultiCarrierEncodeByFileNames(carrierFileNames []string, dataFileName string, resultFileNames []string) (err error) {
130	if len(carrierFileNames) == 0 {
131		return fmt.Errorf("missing carriers names")
132	}
133	if len(carrierFileNames) != len(resultFileNames) {
134		return fmt.Errorf("different number of carriers and results")
135	}
136	carriers := make([]io.Reader, 0, len(carrierFileNames))
137	for _, name := range carrierFileNames {
138		carrier, err := os.Open(name)
139		if err != nil {
140			return fmt.Errorf("error opening carrier file %s: %v", name, err)
141		}
142		defer func() {
143			closeErr := carrier.Close()
144			if err == nil {
145				err = closeErr
146			}
147		}()
148		carriers = append(carriers, carrier)
149	}
150
151	data, err := os.Open(dataFileName)
152	if err != nil {
153		return fmt.Errorf("error opening data file %s: %v", dataFileName, err)
154	}
155	defer func() {
156		closeErr := data.Close()
157		if err == nil {
158			err = closeErr
159		}
160	}()
161
162	results := make([]io.Writer, 0, len(resultFileNames))
163	for _, name := range resultFileNames {
164		result, err := os.Create(name)
165		if err != nil {
166			return fmt.Errorf("error creating result file %s: %v", name, err)
167		}
168		defer func() {
169			closeErr := result.Close()
170			if err == nil {
171				err = closeErr
172			}
173		}()
174		results = append(results, result)
175	}
176
177	err = MultiCarrierEncode(carriers, data, results)
178	if err != nil {
179		for _, name := range resultFileNames {
180			_ = os.Remove(name)
181		}
182	}
183	return err
184}
185
186func quartersOfBytesOf(counter uint32) []byte {
187	bs := make([]byte, 4)
188	binary.LittleEndian.PutUint32(bs, counter)
189	quarters := make([]byte, 16)
190	for i := 0; i < 16; i += 4 {
191		quarters[i] = bits.QuartersOfByte(bs[i/4])[0]
192		quarters[i+1] = bits.QuartersOfByte(bs[i/4])[1]
193		quarters[i+2] = bits.QuartersOfByte(bs[i/4])[2]
194		quarters[i+3] = bits.QuartersOfByte(bs[i/4])[3]
195	}
196
197	return quarters
198}
199
200func setDataSizeHeader(RGBAImage *image.RGBA, dataCountBytes []byte) {
201	dx := RGBAImage.Bounds().Dx()
202	dy := RGBAImage.Bounds().Dy()
203
204	count := 0
205
206	for x := 0; x < dx && count < (dataSizeHeaderReservedBytes/4)*3; x++ {
207		for y := 0; y < dy && count < (dataSizeHeaderReservedBytes/4)*3; y++ {
208			c := RGBAImage.RGBAAt(x, y)
209			c.R = bits.SetLastTwoBits(c.R, dataCountBytes[count])
210			c.G = bits.SetLastTwoBits(c.G, dataCountBytes[count+1])
211			c.B = bits.SetLastTwoBits(c.B, dataCountBytes[count+2])
212			RGBAImage.SetRGBA(x, y, c)
213
214			count += 3
215
216		}
217	}
218}
219
220func setColorSegment(colorSegment *byte, data <-chan byte, errChan <-chan error) (hasMoreBytes bool, err error) {
221	select {
222	case byte, ok := <-data:
223		if !ok {
224			return false, nil
225		}
226		*colorSegment = bits.SetLastTwoBits(*colorSegment, byte)
227		return true, nil
228
229	case err := <-errChan:
230		return false, err
231
232	}
233}
234
235func readData(reader io.Reader, bytes chan<- byte, errChan chan<- error) {
236	b := make([]byte, 1)
237	for {
238		if _, err := reader.Read(b); err != nil {
239			if err == io.EOF {
240				break
241			}
242			errChan <- fmt.Errorf("error reading data %v", err)
243			return
244		}
245		for _, b := range bits.QuartersOfByte(b[0]) {
246			bytes <- b
247		}
248
249	}
250	close(bytes)
251}
252
253func getImageAsRGBA(reader io.Reader) (*image.RGBA, string, error) {
254	img, format, err := image.Decode(reader)
255	if err != nil {
256		return nil, format, fmt.Errorf("error decoding carrier image: %v", err)
257	}
258
259	RGBAImage := image.NewRGBA(image.Rect(0, 0, img.Bounds().Dx(), img.Bounds().Dy()))
260	draw.Draw(RGBAImage, RGBAImage.Bounds(), img, img.Bounds().Min, draw.Src)
261
262	return RGBAImage, format, nil
263}
264