Add JPG and JPEG support

cmd/lsb/main.go

@@ -8,7 +8,9 @@ import (
 	"github.com/pkg/errors"
 	"github.com/urfave/cli/v3"
 	"image"
+	"image/jpeg"
 	"os"
+	"steganography/internal/helper"
 	"steganography/internal/lsb"
 )
 
@@ -123,9 +125,22 @@ func executeEncode(_ context.Context, command *cli.Command) error {
 	}
 	defer inputFile.Close()
 
-	// Reads binary data from input file
+	// Get content type
 	reader := bufio.NewReader(inputFile)
-	img, _, err := image.Decode(reader)
+	bts, _ := reader.Peek(512)
+	contentType := helper.GetContentType(bts)
+
+	// Reads binary data from input file
+	var img image.Image
+	img, _, err = image.Decode(reader)
+	switch {
+	case helper.IsPngContentType(contentType):
+
+	case helper.IsJpgContentType(contentType) || helper.IsJpegContentType(contentType):
+		img, err = jpeg.Decode(reader)
+	default:
+		err = errors.New("unknown image format: image format must be PNG, JPG or JPEG")
+	}
 	if err != nil {
 		return errors.Wrap(err, "failed to decode input image")
 	}
@@ -179,9 +194,22 @@ func executeDecode(_ context.Context, command *cli.Command) error {
 	}
 	defer inputFile.Close()
 
-	// Reads binary data from input file
+	// Get content type
 	reader := bufio.NewReader(inputFile)
-	img, _, err := image.Decode(reader)
+	bts, _ := reader.Peek(512)
+	contentType := helper.GetContentType(bts)
+
+	// Reads binary data from input file
+	var img image.Image
+	img, _, err = image.Decode(reader)
+	switch {
+	case helper.IsPngContentType(contentType):
+
+	case helper.IsJpgContentType(contentType) || helper.IsJpegContentType(contentType):
+		img, err = jpeg.Decode(reader)
+	default:
+		err = errors.New("unknown image format: image format must be PNG, JPG or JPEG")
+	}
 	if err != nil {
 		return errors.Wrap(err, "failed to decode input image")
 	}

internal/helper/image.go

@@ -0,0 +1,28 @@
+package helper
+
+import (
+	"mime"
+	"net/http"
+)
+
+var (
+	pngContentType  = mime.TypeByExtension(".png")
+	jpgContentType  = mime.TypeByExtension(".jpg")
+	jpegContentType = mime.TypeByExtension(".jpeg")
+)
+
+func GetContentType(content []byte) string {
+	return http.DetectContentType(content)
+}
+
+func IsPngContentType(contentType string) bool {
+	return contentType == pngContentType
+}
+
+func IsJpgContentType(contentType string) bool {
+	return contentType == jpgContentType
+}
+
+func IsJpegContentType(contentType string) bool {
+	return contentType == jpegContentType
+}

internal/lsb/steganography_test.go

@@ -4,55 +4,171 @@ import (
 	"bufio"
 	"bytes"
 	"image"
+	"image/jpeg"
 	"log"
 	"os"
 	"steganography/internal/lsb"
 	"testing"
 )
 
 var (
-	testDataPath = getEnvWithDefult("TEST_DATA_DIR", ".")
+	testDataPath = getEnvWithDefault("TEST_DATA_DIR", ".")
 
-	file        = testDataPath + "/statham.png"
-	encodedFile = testDataPath + "/encoded_statham.png"
+	pngFile        = testDataPath + "/statham.png"
+	encodedPngFile = testDataPath + "/encoded_statham.png"
+
+	jpgFile        = testDataPath + "/statham.jpg"
+	encodedJpgFile = testDataPath + "/encoded_statham.jpg"
+
+	jpegFile        = testDataPath + "/statham.jpeg"
+	encodedJpegFile = testDataPath + "/encoded_statham.jpeg"
 
 	message = []byte("Those who get up early want to sleep all day.")
 )
 
-func TestEncodeFromFile(t *testing.T) {
+func TestEncodeFromPngFile(t *testing.T) {
+	inFile, err := os.Open(pngFile)
+	if err != nil {
+		log.Printf("Error opening pngFile %s: %v", pngFile, err)
+		t.FailNow()
+	}
+	defer inFile.Close()
+
+	reader := bufio.NewReader(inFile)
+	img, _, err := image.Decode(reader)
+	if err != nil {
+		log.Printf("Error decoding. %v", err)
+		t.FailNow()
+	}
+	w := new(bytes.Buffer)
+	err = lsb.Encode(w, img, message) // Encode the message into the image file
+	if err != nil {
+		log.Printf("Error encoding file %v", err)
+		t.FailNow()
+	}
+	outFile, err := os.Create(encodedPngFile)
+	if err != nil {
+		log.Printf("Error creating file %s: %v", encodedPngFile, err)
+		t.FailNow()
+	}
+	w.WriteTo(outFile)
+	defer outFile.Close()
+}
+
+func TestDecodeFromPngFile(t *testing.T) {
+	inFile, err := os.Open(encodedPngFile)
+	if err != nil {
+		log.Printf("Error opening file %s: %v", encodedPngFile, err)
+		t.FailNow()
+	}
+	defer inFile.Close()
+
+	reader := bufio.NewReader(inFile)
+	img, _, err := image.Decode(reader)
+	if err != nil {
+		log.Print("Error decoding file")
+		t.FailNow()
+	}
+
+	sizeOfMessage := lsb.GetMessageSizeFromImage(img)
+
+	msg := lsb.Decode(sizeOfMessage, img) // Read the message from the picture pngFile
+
+	if !bytes.Equal(msg, message) {
+		log.Print("messages dont match:")
+		log.Println(string(msg))
+		t.FailNow()
+	}
+}
+
+func TestEncodeFromJpgFile(t *testing.T) {
+	inFile, err := os.Open(jpgFile)
+	if err != nil {
+		log.Printf("Error opening file %s: %v", jpgFile, err)
+		t.FailNow()
+	}
+	defer inFile.Close()
+
+	reader := bufio.NewReader(inFile)
+	img, err := jpeg.Decode(reader)
+	if err != nil {
+		log.Printf("Error decoding. %v", err)
+		t.FailNow()
+	}
+	w := new(bytes.Buffer)
+	err = lsb.Encode(w, img, message) // Encode the message into the image file
+	if err != nil {
+		log.Printf("Error encoding file %v", err)
+		t.FailNow()
+	}
+	outFile, err := os.Create(encodedJpgFile)
+	if err != nil {
+		log.Printf("Error creating file %s: %v", encodedJpgFile, err)
+		t.FailNow()
+	}
+	w.WriteTo(outFile)
+	defer outFile.Close()
+}
 
-	inFile, err := os.Open(file)
+func TestDecodeFromJpgFile(t *testing.T) {
+	inFile, err := os.Open(encodedJpgFile)
 	if err != nil {
-		log.Printf("Error opening file %s: %v", file, err)
+		log.Printf("Error opening file %s: %v", encodedJpgFile, err)
 		t.FailNow()
 	}
 	defer inFile.Close()
 
 	reader := bufio.NewReader(inFile)
 	img, _, err := image.Decode(reader)
+	if err != nil {
+		log.Print("Error decoding file")
+		t.FailNow()
+	}
+
+	sizeOfMessage := lsb.GetMessageSizeFromImage(img)
+
+	msg := lsb.Decode(sizeOfMessage, img) // Read the message from the picture file
+
+	if !bytes.Equal(msg, message) {
+		log.Print("messages dont match:")
+		log.Println(string(msg))
+		t.FailNow()
+	}
+}
+
+func TestEncodeFromJpegFile(t *testing.T) {
+	inFile, err := os.Open(jpegFile)
+	if err != nil {
+		log.Printf("Error opening file %s: %v", jpegFile, err)
+		t.FailNow()
+	}
+	defer inFile.Close()
+
+	reader := bufio.NewReader(inFile)
+	img, err := jpeg.Decode(reader)
 	if err != nil {
 		log.Printf("Error decoding. %v", err)
 		t.FailNow()
 	}
 	w := new(bytes.Buffer)
 	err = lsb.Encode(w, img, message) // Encode the message into the image file
 	if err != nil {
-		log.Printf("Error Encoding file %v", err)
+		log.Printf("Error encoding file %v", err)
 		t.FailNow()
 	}
-	outFile, err := os.Create(encodedFile)
+	outFile, err := os.Create(encodedJpegFile)
 	if err != nil {
-		log.Printf("Error creating file %s: %v", encodedFile, err)
+		log.Printf("Error creating file %s: %v", encodedJpegFile, err)
 		t.FailNow()
 	}
 	w.WriteTo(outFile)
 	defer outFile.Close()
 }
 
-func TestDecodeFromFile(t *testing.T) {
-	inFile, err := os.Open(encodedFile)
+func TestDecodeFromJpegFile(t *testing.T) {
+	inFile, err := os.Open(encodedJpegFile)
 	if err != nil {
-		log.Printf("Error opening file %s: %v", encodedFile, err)
+		log.Printf("Error opening file %s: %v", encodedJpegFile, err)
 		t.FailNow()
 	}
 	defer inFile.Close()
@@ -75,7 +191,7 @@ func TestDecodeFromFile(t *testing.T) {
 	}
 }
 
-func getEnvWithDefult(key, value string) string {
+func getEnvWithDefault(key, value string) string {
 	if val, ok := os.LookupEnv(key); ok {
 		return val
 	}