Add Write

cmd/zwfp/main.go

@@ -1,31 +1,60 @@
 package main
 
 import (
+	"flag"
 	"fmt"
+	"io/ioutil"
+	"log"
 	"os"
 
 	"github.com/vedhavyas/zwfp"
 )
 
 func main() {
-	switch len(os.Args) {
-	case 3:
-		embed := zwfp.Embed(os.Args[1], os.Args[2])
-		fmt.Println(embed)
-		return
-	case 2:
-		pt, key := zwfp.Extract(os.Args[1])
-		fmt.Println("Cover Text:", pt)
-		fmt.Println("Payload:", key)
-		return
-	default:
+	var key string
+	flag.StringVar(&key, "k", "", "key string")
+	flag.Usage = func() {
+		name := os.Args[0]
 		fmt.Println("Usage:")
-		fmt.Println("\t", os.Args[0], "CoverText Payload")
-		fmt.Println("\t\t", "Embeds Payload into CoverText")
-		fmt.Println("")
-		fmt.Println("\t", os.Args[0], "SteganoText")
-		fmt.Println("\t\t", "Extracts Payload from SteganoText")
-		os.Exit(1)
+		flag.PrintDefaults()
+		fmt.Printf(`
+  Embeding:
+    $ cat covertext.txt | %[1]s -k Payload
+    $ %[1]s -k Payload CoverText
+
+  Extracting:
+    $ cat covertext.txt | %[1]s
+    $ %[1]s SteganoText
+`, name)
 	}
+	flag.Parse()
 
+	if key == "" {
+		text := flag.Arg(0)
+		if text == "" {
+			b, err := ioutil.ReadAll(os.Stdin)
+			if err != nil {
+				log.Fatal(err)
+			}
+			_, key := zwfp.Extract(string(b))
+			fmt.Println(key)
+		} else {
+			pt, key := zwfp.Extract(text)
+			fmt.Println("Cover Text:", pt)
+			fmt.Println("Payload:", key)
+		}
+	} else {
+		switch flag.NArg() {
+		case 0:
+			if err := zwfp.Write(os.Stdout, os.Stdin, key); err != nil {
+				log.Fatal(err)
+			}
+		case 1:
+			embed := zwfp.Embed(flag.Arg(0), key)
+			fmt.Println(embed)
+		default:
+			flag.Usage()
+			os.Exit(1)
+		}
+	}
 }

embed.go

@@ -1,7 +1,9 @@
 package zwfp
 
 import (
+	"bufio"
 	"fmt"
+	"io"
 	"strings"
 )
 
@@ -108,3 +110,78 @@ func Embed(data, key string) string {
 
 	return string(embed)
 }
+
+// Write embeds the key into stream Writer from Reader using zero-width characters
+func Write(w io.Writer, r io.Reader, key string) error {
+	zwKey := toZeroWidth(key)
+	var zwRKey []rune
+	for _, c := range zwKey {
+		zwRKey = append(zwRKey, c)
+	}
+
+	var t int
+	var err error
+	var c rune
+
+	i := 0
+	br := bufio.NewReader(r)
+	bw := bufio.NewWriter(w)
+	defer bw.Flush()
+	for {
+		c, _, err = br.ReadRune()
+		if err != nil {
+			break
+		}
+		if i == 0 {
+			_, err = bw.WriteRune(c)
+			if err != nil {
+				break
+			}
+		}
+
+		if t < len(zwRKey) {
+			_, err = bw.WriteRune(zwRKey[t])
+			if err != nil {
+				break
+			}
+			t++
+		}
+
+		if i != 0 {
+			_, err = bw.WriteRune(c)
+			if err != nil {
+				break
+			}
+		}
+		i++
+	}
+
+	if err == io.EOF {
+		err = nil
+	}
+
+	if t < len(zwRKey) {
+		if bw.Size() > 0 {
+			br.UnreadByte()
+			_, err = bw.WriteString(string(zwRKey[t:]))
+			if err != nil {
+				return err
+			}
+			c, _, err = br.ReadRune()
+			if err != nil {
+				return err
+			}
+			_, err = bw.WriteRune(c)
+			if err != nil {
+				return err
+			}
+		} else {
+			_, err = bw.WriteString(string(zwRKey[t:]))
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return err
+}