From 2509420ba504cde98260f47b02c0165da6d6ee7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ken-H=C3=A5vard=20Lieng?= Date: Tue, 12 May 2020 04:13:05 +0200 Subject: [PATCH] Make ToCTCP handle any type of CTCP message, move DCC handling to separate file --- pkg/irc/client.go | 152 ------------------------------ pkg/irc/conn.go | 20 +++- pkg/irc/dcc.go | 227 +++++++++++++++++++++++++++++++++++++++++++++ pkg/irc/message.go | 83 ++++------------- 4 files changed, 261 insertions(+), 221 deletions(-) create mode 100644 pkg/irc/dcc.go diff --git a/pkg/irc/client.go b/pkg/irc/client.go index e3d2ed38..83dd40b0 100644 --- a/pkg/irc/client.go +++ b/pkg/irc/client.go @@ -3,13 +3,7 @@ package irc import ( "bufio" "crypto/tls" - "encoding/binary" - "fmt" - "io" - "math" "net" - "os" - "path/filepath" "strings" "sync" "time" @@ -218,149 +212,3 @@ func (c *Client) flushChannels() { } c.lock.Unlock() } - -func byteRead(totalBytes uint64) []byte { - b := make([]byte, 8) - binary.BigEndian.PutUint64(b, totalBytes) - return b -} - -func round2(source float64) float64 { - return math.Round(100*source) / 100 -} - -const ( - _ = 1.0 << (10 * iota) - kibibyte - mebibyte - gibibyte -) - -func humanReadableByteCount(b float64, speed bool) string { - unit := "" - value := b - - switch { - case b >= gibibyte: - unit = "GiB" - value = value / gibibyte - case b >= mebibyte: - unit = "MiB" - value = value / mebibyte - case b >= kibibyte: - unit = "KiB" - value = value / kibibyte - case b > 1 || b == 0: - unit = "bytes" - case b == 1: - unit = "byte" - } - - if speed { - unit = unit + "/s" - } - - stringValue := strings.TrimSuffix( - fmt.Sprintf("%.2f", value), ".00", - ) - - return fmt.Sprintf("%s %s", stringValue, unit) -} - -func (c *Client) Download(pack *CTCP) { - if !c.Autoget { - // TODO: ask user if he/she wants to download the file - return - } - c.Progress <- DownloadProgress{ - PercCompletion: 0, - File: pack.File, - } - file, err := os.OpenFile(filepath.Join(c.DownloadFolder, pack.File), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) - if err != nil { - c.Progress <- DownloadProgress{ - PercCompletion: -1, - File: pack.File, - Error: err, - } - return - } - defer file.Close() - - con, err := net.Dial("tcp", fmt.Sprintf("%s:%d", pack.IP, pack.Port)) - - if err != nil { - c.Progress <- DownloadProgress{ - PercCompletion: -1, - File: pack.File, - Error: err, - } - return - } - - defer con.Close() - - var avgSpeed float64 - var prevTime int64 = -1 - secondsElapsed := int64(0) - totalBytes := uint64(0) - buf := make([]byte, 0, 4*1024) - start := time.Now().UnixNano() - for { - n, err := con.Read(buf[:cap(buf)]) - buf = buf[:n] - if n == 0 { - if err == nil { - continue - } - if err == io.EOF { - break - } - } - - if _, err := file.Write(buf); err != nil { - c.Progress <- DownloadProgress{ - PercCompletion: -1, - File: pack.File, - Error: err, - } - return - } - - cycleBytes := uint64(len(buf)) - totalBytes += cycleBytes - percentage := round2(100 * float64(totalBytes) / float64(pack.Length)) - - now := time.Now().UnixNano() - secondsElapsed = (now - start) / 1e9 - avgSpeed = round2(float64(totalBytes) / (float64(secondsElapsed))) - speed := 0.0 - - if prevTime < 0 { - speed = avgSpeed - } else { - speed = round2(1e9 * float64(cycleBytes) / (float64(now - prevTime))) - } - secondsToGo := (float64(pack.Length) - float64(totalBytes)) / speed - prevTime = now - con.Write(byteRead(totalBytes)) - c.Progress <- DownloadProgress{ - InstSpeed: humanReadableByteCount(speed, true), - AvgSpeed: humanReadableByteCount(avgSpeed, true), - PercCompletion: percentage, - BytesRemaining: humanReadableByteCount(float64(pack.Length-totalBytes), false), - BytesCompleted: humanReadableByteCount(float64(totalBytes), false), - SecondsElapsed: secondsElapsed, - SecondsToGo: secondsToGo, - File: pack.File, - } - } - con.Write(byteRead(totalBytes)) - c.Progress <- DownloadProgress{ - AvgSpeed: humanReadableByteCount(avgSpeed, true), - PercCompletion: 100, - BytesCompleted: humanReadableByteCount(float64(totalBytes), false), - SecondsElapsed: secondsElapsed, - File: pack.File, - } -} diff --git a/pkg/irc/conn.go b/pkg/irc/conn.go index ae819a2b..fc8b698a 100644 --- a/pkg/irc/conn.go +++ b/pkg/irc/conn.go @@ -223,6 +223,11 @@ func (c *Client) recv() { c.setNick(msg.LastParam()) } + case Privmsg: + if ctcp := msg.ToCTCP(); ctcp != nil { + c.handleCTCP(ctcp) + } + case ReplyWelcome: c.setNick(msg.Params[0]) c.setRegistered(true) @@ -242,10 +247,17 @@ func (c *Client) recv() { } - if ctcp := msg.ToCTCP(); ctcp != nil { - go c.Download(ctcp) - } - c.Messages <- msg } } + +func (c *Client) handleCTCP(ctcp *CTCP) { + switch ctcp.Command { + case "DCC": + if strings.HasPrefix(ctcp.Params, "SEND") { + if dccSend := ParseDCCSend(ctcp); dccSend != nil { + go c.Download(dccSend) + } + } + } +} diff --git a/pkg/irc/dcc.go b/pkg/irc/dcc.go new file mode 100644 index 00000000..014a24c2 --- /dev/null +++ b/pkg/irc/dcc.go @@ -0,0 +1,227 @@ +package irc + +import ( + "encoding/binary" + "encoding/json" + "fmt" + "io" + "math" + "net" + "os" + "path" + "path/filepath" + "strconv" + "strings" + "time" +) + +type DCCSend struct { + File string `json:"file"` + IP string `json:"ip"` + Port uint16 `json:"port"` + Length uint64 `json:"length"` +} + +func ParseDCCSend(ctcp *CTCP) *DCCSend { + params := strings.Split(ctcp.Params, " ") + + if len(params) > 4 { + ip, err := strconv.Atoi(params[2]) + port, err := strconv.Atoi(params[3]) + length, err := strconv.Atoi(params[4]) + + if err != nil { + return nil + } + + ip3 := uint32ToIP(ip) + + filename := path.Base(params[1]) + if filename == "/" || filename == "." { + filename = "" + } + + return &DCCSend{ + File: filename, + IP: ip3, + Port: uint16(port), + Length: uint64(length), + } + } + + return nil +} + +func (c *Client) Download(pack *DCCSend) { + if !c.Autoget { + // TODO: ask user if he/she wants to download the file + return + } + c.Progress <- DownloadProgress{ + PercCompletion: 0, + File: pack.File, + } + file, err := os.OpenFile(filepath.Join(c.DownloadFolder, pack.File), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + c.Progress <- DownloadProgress{ + PercCompletion: -1, + File: pack.File, + Error: err, + } + return + } + defer file.Close() + + con, err := net.Dial("tcp", fmt.Sprintf("%s:%d", pack.IP, pack.Port)) + + if err != nil { + c.Progress <- DownloadProgress{ + PercCompletion: -1, + File: pack.File, + Error: err, + } + return + } + + defer con.Close() + + var avgSpeed float64 + var prevTime int64 = -1 + secondsElapsed := int64(0) + totalBytes := uint64(0) + buf := make([]byte, 0, 4*1024) + start := time.Now().UnixNano() + for { + n, err := con.Read(buf[:cap(buf)]) + buf = buf[:n] + if n == 0 { + if err == nil { + continue + } + if err == io.EOF { + break + } + } + + if _, err := file.Write(buf); err != nil { + c.Progress <- DownloadProgress{ + PercCompletion: -1, + File: pack.File, + Error: err, + } + return + } + + cycleBytes := uint64(len(buf)) + totalBytes += cycleBytes + percentage := round2(100 * float64(totalBytes) / float64(pack.Length)) + + now := time.Now().UnixNano() + secondsElapsed = (now - start) / 1e9 + avgSpeed = round2(float64(totalBytes) / (float64(secondsElapsed))) + speed := 0.0 + + if prevTime < 0 { + speed = avgSpeed + } else { + speed = round2(1e9 * float64(cycleBytes) / (float64(now - prevTime))) + } + secondsToGo := (float64(pack.Length) - float64(totalBytes)) / speed + prevTime = now + con.Write(byteRead(totalBytes)) + c.Progress <- DownloadProgress{ + InstSpeed: humanReadableByteCount(speed, true), + AvgSpeed: humanReadableByteCount(avgSpeed, true), + PercCompletion: percentage, + BytesRemaining: humanReadableByteCount(float64(pack.Length-totalBytes), false), + BytesCompleted: humanReadableByteCount(float64(totalBytes), false), + SecondsElapsed: secondsElapsed, + SecondsToGo: secondsToGo, + File: pack.File, + } + } + con.Write(byteRead(totalBytes)) + c.Progress <- DownloadProgress{ + AvgSpeed: humanReadableByteCount(avgSpeed, true), + PercCompletion: 100, + BytesCompleted: humanReadableByteCount(float64(totalBytes), false), + SecondsElapsed: secondsElapsed, + File: pack.File, + } +} + +type DownloadProgress struct { + File string `json:"file"` + Error error `json:"error"` + BytesCompleted string `json:"bytes_completed"` + BytesRemaining string `json:"bytes_remaining"` + PercCompletion float64 `json:"perc_completion"` + AvgSpeed string `json:"avg_speed"` + InstSpeed string `json:"speed"` + SecondsElapsed int64 `json:"elapsed"` + SecondsToGo float64 `json:"eta"` +} + +func (p DownloadProgress) ToJSON() string { + progress, err := json.Marshal(p) + if err != nil { + return "" + } + return string(progress) +} + +func uint32ToIP(n int) string { + var byte1 = n & 255 + var byte2 = ((n >> 8) & 255) + var byte3 = ((n >> 16) & 255) + var byte4 = ((n >> 24) & 255) + return fmt.Sprintf("%d.%d.%d.%d", byte4, byte3, byte2, byte1) +} + +func byteRead(totalBytes uint64) []byte { + b := make([]byte, 8) + binary.BigEndian.PutUint64(b, totalBytes) + return b +} + +func round2(source float64) float64 { + return math.Round(100*source) / 100 +} + +const ( + _ = 1.0 << (10 * iota) + kibibyte + mebibyte + gibibyte +) + +func humanReadableByteCount(b float64, speed bool) string { + unit := "" + value := b + + switch { + case b >= gibibyte: + unit = "GiB" + value = value / gibibyte + case b >= mebibyte: + unit = "MiB" + value = value / mebibyte + case b >= kibibyte: + unit = "KiB" + value = value / kibibyte + case b > 1 || b == 0: + unit = "bytes" + case b == 1: + unit = "byte" + } + + if speed { + unit = unit + "/s" + } + + stringValue := strings.TrimSuffix( + fmt.Sprintf("%.2f", value), ".00", + ) + + return fmt.Sprintf("%s %s", stringValue, unit) +} diff --git a/pkg/irc/message.go b/pkg/irc/message.go index b40fecd7..6327fec4 100644 --- a/pkg/irc/message.go +++ b/pkg/irc/message.go @@ -1,12 +1,7 @@ package irc import ( - "encoding/json" - "fmt" - "path" - "strconv" "strings" - "unicode" ) type Message struct { @@ -17,26 +12,6 @@ type Message struct { Params []string } -type DownloadProgress struct { - File string `json:"file"` - Error error `json:"error"` - BytesCompleted string `json:"bytes_completed"` - BytesRemaining string `json:"bytes_remaining"` - PercCompletion float64 `json:"perc_completion"` - AvgSpeed string `json:"avg_speed"` - InstSpeed string `json:"speed"` - SecondsElapsed int64 `json:"elapsed"` - SecondsToGo float64 `json:"eta"` -} - -// CTCP is used to parse a message into a CTCP message -type CTCP struct { - File string `json:"file"` - IP string `json:"ip"` - Port uint16 `json:"port"` - Length uint64 `json:"length"` -} - func (m *Message) LastParam() string { if len(m.Params) > 0 { return m.Params[len(m.Params)-1] @@ -44,37 +19,31 @@ func (m *Message) LastParam() string { return "" } -// ToCTCP tries to parse the message parameters into a CTCP message -func (m *Message) ToCTCP() *CTCP { - params := strings.Join(m.Params, " ") - if strings.Contains(params, "DCC SEND") { - // to be extra sure that there are non-printable characters - params = strings.TrimFunc(params, func(r rune) bool { - return !unicode.IsPrint(r) - }) - parts := strings.Split(params, " ") - ip, err := strconv.Atoi(parts[4]) - port, err := strconv.Atoi(parts[5]) - length, err := strconv.Atoi(parts[6]) +type CTCP struct { + Command string + Params string +} - if err != nil { +func (m *Message) ToCTCP() *CTCP { + lp := m.LastParam() + + if len(lp) > 1 && lp[0] == 0x01 { + parts := strings.SplitN(strings.Trim(lp, "\x01"), " ", 2) + ctcp := CTCP{} + + if parts[0] != "" { + ctcp.Command = parts[0] + } else { return nil } - ip3 := uint32ToIP(ip) - - filename := path.Base(parts[3]) - if filename == "/" || filename == "." { - filename = "" + if len(parts) == 2 { + ctcp.Params = parts[1] } - return &CTCP{ - File: filename, - IP: ip3, - Port: uint16(port), - Length: uint64(length), - } + return &ctcp } + return nil } @@ -171,19 +140,3 @@ var unescapeTagReplacer = strings.NewReplacer( func unescapeTag(s string) string { return unescapeTagReplacer.Replace(s) } - -func uint32ToIP(n int) string { - var byte1 = n & 255 - var byte2 = ((n >> 8) & 255) - var byte3 = ((n >> 16) & 255) - var byte4 = ((n >> 24) & 255) - return fmt.Sprintf("%d.%d.%d.%d", byte4, byte3, byte2, byte1) -} - -func (p DownloadProgress) ToJSON() string { - progress, err := json.Marshal(p) - if err != nil { - return "" - } - return string(progress) -}