2 files changed, 22 insertions, 12 deletions
diff --git a/stdin-service/default.txt b/stdin-service/default.txt
index 78ed276..7518f53 100644
--- a/stdin-service/default.txt
+++ b/stdin-service/default.txt
@@ -998,4 +998,3 @@
 997
 998
 999
-
diff --git a/stdin-service/main.go b/stdin-service/main.go
index 39843a1..df0c503 100644
--- a/stdin-service/main.go
+++ b/stdin-service/main.go
@@ -3,6 +3,7 @@ package main
 import (
     "flag"
     "io/ioutil"
+    "io"
     "log"
     "os"
     "net"
@@ -14,25 +15,35 @@ var (
     host = flag.String("h", "localhost", "the host to pee against")
     port = flag.Int("p", 1337, "the port to use")
     protocol = flag.String("protocol", "tcp", "the protocol to use")
+    stdin = flag.Bool("stdin", false, "read from stdin")
 )
 
 func main() {
     flag.Parse()
 
-    dat, err := ioutil.ReadFile(*inputFile)
-    if err != nil {
-        log.Printf("Error reading file: %v", err)
-        os.Exit(1)
-    }
-
-    conn, _ := net.Dial(*protocol, fmt.Sprintf("%s:%d", *host, *port))
+    var dat []byte
 
-    writtenBytes, err := conn.Write(dat)
-    if err != nil {
-            log.Printf("Error accepting: %v", err)
+    if !*stdin {
+        var ReadFileErr error
+        dat, ReadFileErr = ioutil.ReadFile(*inputFile)
+        if ReadFileErr != nil {
+            log.Printf("Error reading file: %v", ReadFileErr)
             os.Exit(1)
+        }
+    } else {
+        dat = make([]byte, 1000)
+        _, err := io.ReadFull(os.Stdin, dat)
+        if err != nil {
+            log.Printf("Error reading from stdin: %v", err)
+        }
     }
-    fmt.Printf("Written %d bytes", writtenBytes)
 
+   conn, _ := net.Dial(*protocol, fmt.Sprintf("%s:%d", *host, *port))
 
+   writtenBytes, err := conn.Write(dat)
+   if err != nil {
+           log.Printf("Error accepting: %v", err)
+           os.Exit(1)
+   }
+   fmt.Printf("Written %d bytes", writtenBytes)
 }