From 7b1460a5fc74b5c1ef9e9e459509732047aeae42 Mon Sep 17 00:00:00 2001
From: Dantezy <zhangyet@gmail.com>
Date: Thu, 4 Jul 2024 00:10:00 +0800
Subject: [PATCH] Add ofile option.

Signed-off-by: Dantezy <zhangyet@gmail.com>
---
 README.md               |  3 ++-
 internal/dust/output.go | 12 +++++++++---
 internal/dust/types.go  |  2 ++
 main.go                 |  2 +-
 4 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 3df06c4..5d38eaa 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,7 @@ Usage: ./dust [options]
     Available options:
   -h, --help              show help
   -i, --interval uint32   set monitor time in seconds (default 10)
+  -o, --ofile string      output file
   -p, --pid uint32        filter pid
   -v, --version           show version
 ```
@@ -24,4 +25,4 @@ sudo ./dust -p {the_process_you_want_to_trace} -i 30
 ## TODO
 - [ ] kprobe filters.
 - [ ] kprobe.multi supports.
-- [ ] output to files.
\ No newline at end of file
+- [x] output to files.
\ No newline at end of file
diff --git a/internal/dust/output.go b/internal/dust/output.go
index 86410f6..f708ab2 100644
--- a/internal/dust/output.go
+++ b/internal/dust/output.go
@@ -20,6 +20,7 @@ type output struct {
 	kprobeMulti  bool // TODO
 	kfreeReasons map[uint64]string
 	ifaceCache   map[uint64]map[uint32]string
+	ofile        string
 }
 
 // outputStructured is a struct to hold the data for the json output
@@ -48,8 +49,15 @@ type jsonTuple struct {
 	Proto uint8  `json:"proto,omitempty"`
 }
 
-func NewOutput(addr2Name Addr2Name, kprobeMulti bool) (*output, error) {
+func NewOutput(addr2Name Addr2Name, kprobeMulti bool, ofile string) (*output, error) {
 	writer := os.Stdout
+	if ofile != "" {
+		var err error
+		writer, err = os.OpenFile(ofile, os.O_RDWR|os.O_CREATE, 0644)
+		if err != nil {
+			return nil, err
+		}
+	}
 
 	var ifs map[uint64]map[uint32]string
 
@@ -70,8 +78,6 @@ func (o *output) Close() {
 }
 
 func (o *output) PrintHeader() {
-	// TODO
-
 }
 
 // PrintJson prints the event in JSON format
diff --git a/internal/dust/types.go b/internal/dust/types.go
index 109877c..5ad89ef 100644
--- a/internal/dust/types.go
+++ b/internal/dust/types.go
@@ -21,6 +21,7 @@ type Flags struct {
 
 	Pid      uint32
 	Interval uint32
+	Ofile    string
 }
 
 func (f *Flags) SetFlags() {
@@ -28,6 +29,7 @@ func (f *Flags) SetFlags() {
 	flag.BoolVarP(&f.ShowHelp, "help", "h", false, "show help")
 	flag.Uint32VarP(&f.Interval, "interval", "i", 10, "set monitor time in seconds")
 	flag.Uint32VarP(&f.Pid, "pid", "p", 0, "filter pid")
+	flag.StringVarP(&f.Ofile, "ofile", "o", "", "output file")
 
 	flag.Usage = func() {
 		_, _ = fmt.Fprintf(os.Stderr, "Usage: %s [options] \n", os.Args[0])
diff --git a/main.go b/main.go
index dd0092e..759e7f1 100644
--- a/main.go
+++ b/main.go
@@ -143,7 +143,7 @@ func main() {
 
 	log.Println("Listening for events..")
 
-	output, err := dust.NewOutput(addr2name, false)
+	output, err := dust.NewOutput(addr2name, false, flags.Ofile)
 	if err != nil {
 		log.Fatalf("Failed to create output: %s", err)
 	}