新增出入口关联挖掘功能

6 files changed, 144 insertions, 19 deletions

diff --git a/Ohmyfile b/Ohmyfile
index 2d8155a..5f72c7a 100644
--- a/Ohmyfile
+++ b/Ohmyfile
@@ -1,20 +1,4 @@
-#.:53 {
-#    v64dns n64.top {
-#        v4ns ns4-1 8.217.120.200
-#        v6ns ns6-1 240b:4001:21b:d301:951b:8de:3cd7:c0e5
-#        chain v4-1 v6-1
-#        ip-embed
-#        chain-mlen 4
-#        analyze neo4j http://localhost:7447 neo4j hdk19990815
-#    }
-#}
-
-
-#prober:2115 {
-#    qname
-#    probe53
-#}
-
 .:53 {
-    atk comm.a.net ns.a.net 1.2.3.4 v6.b.net ns.b.net fe80::6cc1:aafd:5039:7b2b 10
+    debug
+    rtest ns6-1.n64.top. 2001::1 from_to.csv
 }
\ No newline at end of file
diff --git a/core/dnsserver/zdirectives.go b/core/dnsserver/zdirectives.go
index 6844d1f..ba08d92 100644
--- a/core/dnsserver/zdirectives.go
+++ b/core/dnsserver/zdirectives.go
@@ -19,4 +19,5 @@ var Directives = []string{
 	"v64dns",
 	"whoami",
 	"atk",
+	"rtest",
 }
diff --git a/core/plug/zplugin.go b/core/plug/zplugin.go
index 13c77c0..2892482 100644
--- a/core/plug/zplugin.go
+++ b/core/plug/zplugin.go
@@ -13,6 +13,7 @@ import (
 	_ "ohmydns2/plugin/prober/probe53"
 	_ "ohmydns2/plugin/prober/qname"
 	_ "ohmydns2/plugin/prometheus"
+	_ "ohmydns2/plugin/rtest"
 	_ "ohmydns2/plugin/v64dns"
 	_ "ohmydns2/plugin/whoami"
 )
diff --git a/plugin.cfg b/plugin.cfg
index 337fdcf..e5f3213 100644
--- a/plugin.cfg
+++ b/plugin.cfg
@@ -8,4 +8,5 @@ v64dns:v64dns
 whoami:whoami
 qname:qname
 probe53:probe53
-atk:atk
\ No newline at end of file
+atk:atk
+rtest:rtest
\ No newline at end of file
diff --git a/plugin/rtest/rtest.go b/plugin/rtest/rtest.go
new file mode 100644
index 0000000..011b2db
--- /dev/null
+++ b/plugin/rtest/rtest.go
@@ -0,0 +1,110 @@
+package rtest
+
+import (
+	"context"
+	"encoding/csv"
+	"fmt"
+	"github.com/miekg/dns"
+	"net"
+	"ohmydns2/plugin/pkg/request"
+	"os"
+	"strings"
+	"time"
+)
+
+type Rtest struct {
+	addr     string
+	ns       string
+	datachan chan []string
+}
+
+func (r Rtest) ServeDNS(ctx context.Context, writer dns.ResponseWriter, msg *dns.Msg) (int, error) {
+	state := request.Request{W: writer, Req: msg}
+	m := new(dns.Msg)
+	m.SetReply(msg)
+	m.Authoritative = true
+	// 应对0x20
+	q := strings.ToLower(state.QName())
+	qlist := strings.Split(q, ".")
+	// 符合要求的域名
+	if len(qlist) >= 4 {
+		// 1.入口
+		_ = strings.ReplaceAll(qlist[len(qlist)-5], "-", ":")
+		// 2.出口
+		_ = state.IP()
+		// 3.记录关联
+		r.writeData([]string{"20011", "127001"})
+		// 4.返回响应
+		m = r.response(m, state)
+		err := writer.WriteMsg(m)
+		if err != nil {
+			log.Info(err.Error())
+			return dns.RcodeServerFailure, err
+		}
+	}
+	return 0, nil
+}
+
+func (r Rtest) Name() string {
+	return "rtest"
+}
+
+func (r Rtest) newCSVwriterAndWrite(path string) chan []string {
+	f, err := os.OpenFile("./"+time.Now().Format("2006_01_02-15_03_04")+"_"+path, os.O_CREATE|os.O_APPEND|os.O_RDWR, 0666)
+	if err != nil {
+		fmt.Println("Error: ", err)
+		return nil
+	}
+
+	csvWriter := *csv.NewWriter(f)
+	csvWriter.UseCRLF = true
+	var header = []string{"from", "to"}
+	err = csvWriter.Write(header)
+	if err != nil {
+		return nil
+	}
+	csvWriter.Flush()
+	datachan := make(chan []string, 10)
+	go func() {
+		defer func(f *os.File) {
+			err = f.Close()
+			if err != nil {
+				log.Errorf("%v", err.Error())
+			}
+		}(f)
+		for i := range datachan {
+			log.Debug(i)
+			err = csvWriter.Write(i)
+			if err != nil {
+				log.Errorf("%v", err.Error())
+				return
+			}
+			csvWriter.Flush()
+		}
+	}()
+	return datachan
+}
+
+func (r Rtest) writeData(data []string) {
+	r.datachan <- data
+}
+
+func (r Rtest) response(msg *dns.Msg, state request.Request) *dns.Msg {
+	//返回AAAA记录
+	msg.Rcode = dns.RcodeSuccess
+	aaaa := new(dns.AAAA)
+	aaaa.Hdr = dns.RR_Header{Name: state.QName(), Class: dns.ClassINET, Ttl: 10, Rrtype: dns.TypeAAAA}
+	aaaa.AAAA = net.ParseIP(r.addr)
+	msg.Answer = append(msg.Answer, aaaa)
+	//授权记录
+	rec := new(dns.NS)
+	rec.Hdr = dns.RR_Header{Name: "rtest.n64.top.", Class: dns.ClassINET, Ttl: 10, Rrtype: dns.TypeNS}
+	rec.Ns = r.ns
+	msg.Ns = append(msg.Ns, rec)
+	//胶水记录
+	recaddr := new(dns.AAAA)
+	recaddr.Hdr = dns.RR_Header{Name: r.ns, Class: dns.ClassINET, Ttl: 10, Rrtype: dns.TypeAAAA}
+	recaddr.AAAA = net.ParseIP(state.LocalIP())
+	msg.Extra = append(msg.Extra, recaddr)
+	return msg
+}
diff --git a/plugin/rtest/setup.go b/plugin/rtest/setup.go
new file mode 100644
index 0000000..8e7c1f0
--- /dev/null
+++ b/plugin/rtest/setup.go
@@ -0,0 +1,28 @@
+package rtest
+
+import (
+	"github.com/coredns/caddy"
+	"ohmydns2/core/dnsserver"
+	"ohmydns2/plugin"
+	log2 "ohmydns2/plugin/pkg/log"
+)
+
+func init() { plugin.Register("rtest", setup) }
+
+func setup(c *caddy.Controller) error {
+	rt := new(Rtest)
+	c.Next()
+
+	args := c.RemainingArgs()
+	rt.addr = args[1]
+	rt.ns = args[0]
+	rt.datachan = rt.newCSVwriterAndWrite(args[2])
+
+	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
+		return rt
+	})
+
+	return nil
+}
+
+var log = log2.NewWithPlugin("rtest")