diff options
| author | Christine Dodrill <me@christine.website> | 2016-12-09 00:53:57 -0800 |
|---|---|---|
| committer | Christine Dodrill <me@christine.website> | 2016-12-09 00:53:57 -0800 |
| commit | 9eef0e7863920c60983a9a769466ce186ce87661 (patch) | |
| tree | 7a40f52301129ef201697bbee052eb14bcd5ef6b | |
| parent | 496e77598e825fd3958c02b91b6b4f5914e627b0 (diff) | |
| download | x-9eef0e7863920c60983a9a769466ce186ce87661.tar.xz x-9eef0e7863920c60983a9a769466ce186ce87661.zip | |
add simple bnc admin drone
| -rw-r--r-- | irc/bncadmin/.gitignore | 1 | ||||
| -rw-r--r-- | irc/bncadmin/box.rb | 13 | ||||
| -rw-r--r-- | irc/bncadmin/main.go | 148 | ||||
| -rw-r--r-- | irc/bncadmin/run.sh | 11 | ||||
| -rw-r--r-- | irc/bncadmin/vendor-log | 3 | ||||
| -rw-r--r-- | irc/bncadmin/vendor/github.com/belak/irc/client.go | 109 | ||||
| -rw-r--r-- | irc/bncadmin/vendor/github.com/belak/irc/conn.go | 79 | ||||
| -rw-r--r-- | irc/bncadmin/vendor/github.com/belak/irc/handler.go | 16 | ||||
| -rw-r--r-- | irc/bncadmin/vendor/github.com/belak/irc/parser.go | 401 | ||||
| -rw-r--r-- | irc/bncadmin/vendor/github.com/joho/godotenv/autoload/autoload.go | 15 | ||||
| -rw-r--r-- | irc/bncadmin/vendor/github.com/joho/godotenv/godotenv.go | 229 |
11 files changed, 1025 insertions, 0 deletions
diff --git a/irc/bncadmin/.gitignore b/irc/bncadmin/.gitignore new file mode 100644 index 0000000..2eea525 --- /dev/null +++ b/irc/bncadmin/.gitignore @@ -0,0 +1 @@ +.env
\ No newline at end of file diff --git a/irc/bncadmin/box.rb b/irc/bncadmin/box.rb new file mode 100644 index 0000000..fdbc64c --- /dev/null +++ b/irc/bncadmin/box.rb @@ -0,0 +1,13 @@ +from "xena/go" + +workdir "/" +copy "main.go", "/go/src/github.com/Xe/tools/irc/bncadmin/main.go" +copy "vendor", "/go/src/github.com/Xe/tools/irc/bncadmin" +run "go install github.com/Xe/tools/irc/bncadmin && cp /go/bin/bncadmin /usr/bin/bncadmin" + +run "rm -rf /usr/local/go /go && apk del bash gcc musl-dev openssl go" +flatten + +cmd "/usr/bin/bncadmin" + +tag "xena/bncadmin" diff --git a/irc/bncadmin/main.go b/irc/bncadmin/main.go new file mode 100644 index 0000000..99a9b81 --- /dev/null +++ b/irc/bncadmin/main.go @@ -0,0 +1,148 @@ +package main + +import ( + "crypto/tls" + "fmt" + "log" + "os" + "strings" + "sync" + "time" + + "github.com/belak/irc" + _ "github.com/joho/godotenv/autoload" +) + +var ( + bncUsername = needEnv("BNC_USERNAME") + bncPassword = needEnv("BNC_PASSWORD") + bncServer = needEnv("BNC_SERVER") + serverSuffixExpected = needEnv("SERVER_SUFFIX") +) + +func needEnv(key string) string { + v := os.Getenv(key) + if v == "" { + log.Fatal("need value for " + key) + } + + return v +} + +func main() { + log.Println("Bot connecting to " + bncServer) + conn, err := tls.Dial("tcp", bncServer, &tls.Config{ + InsecureSkipVerify: true, + }) + if err != nil { + log.Fatal(err) + } + defer conn.Close() + + c := irc.NewClient(conn, irc.ClientConfig{ + Nick: "admin", + Pass: fmt.Sprintf("%s:%s", bncUsername, bncPassword), + User: "BNCbot", + Name: "BNC admin bot", + + Handler: NewBot(), + }) + + for _, cap := range []string{"userhost-in-names", "multi-prefix", "znc.in/server-time-iso"} { + c.Writef("CAP REQ %s", cap) + } + + err = c.Run() + if err != nil { + panic(err) + } +} + +type Bot struct { + setupDaemon sync.Once + lookingForUserNetworks bool + + // i am sorry + launUsername string +} + +func NewBot() *Bot { + return &Bot{} +} + +func (b *Bot) Handle(c *irc.Client, m *irc.Message) { + b.setupDaemon.Do(func() { + go func() { + b.lookingForUserNetworks = true + c.Writef("PRIVMSG *status ListAllUserNetworks") + time.Sleep(2 * time.Second) // always sleep 2 + b.lookingForUserNetworks = false + + time.Sleep(1 * time.Hour) + }() + }) + + // log.Printf("in >> %s", m) + + switch m.Command { + case "PRIVMSG": + if m.Prefix.Name == "*status" { + b.HandleStarStatus(c, m) + } + + if strings.HasPrefix(m.Prefix.Name, "?") { + b.HandlePartyLineCommand(c, m) + } + + if m.Params[0] == "#bnc" { + b.HandleCommand(c, m) + } + + case "NOTICE": + if m.Prefix.Name == "*status" { + f := strings.Fields(m.Trailing()) + if f[0] == "***" { + log.Println(m.Trailing()) + // look up geoip and log here + } + } + } +} + +func (b *Bot) HandleStarStatus(c *irc.Client, m *irc.Message) { + if b.lookingForUserNetworks { + if strings.HasPrefix(m.Trailing(), "| ") { + f := strings.Fields(m.Trailing()) + + switch len(f) { + case 11: // user name line + // 11: []string{"|", "AzureDiamond", "|", "N/A", "|", "0", "|", "|", "|", "|", "|"} + username := f[1] + b.launUsername = username + + case 15: // server and nick!user@host line + // 15: []string{"|", "`-", "|", "PonyChat", "|", "0", "|", "Yes", "|", "amethyststar.ponychat.net", "|", "test!test@lypbmzxixk.ponychat.net", "|", "1", "|"} + server := f[9] + network := f[3] + if !strings.HasSuffix(server, serverSuffixExpected) { + log.Printf("%s is using the BNC to connect to unknown server %s, removing permissions", b.launUsername, server) + b.RemoveNetwork(c, b.launUsername, network) + c.Writef("PRIVMSG ?%s :You have violated the terms of the BNC service and your account has been disabled. Please contact PonyChat staff to appeal this.", b.launUsername) + c.Writef("PRIVMSG *blockuser block %s", b.launUsername) + } + } + } + } +} + +func (b *Bot) HandlePartyLineCommand(c *irc.Client, m *irc.Message) { + +} + +func (b *Bot) HandleCommand(c *irc.Client, m *irc.Message) { + +} + +func (b *Bot) RemoveNetwork(c *irc.Client, username, network string) { + c.Writef("PRIVMSG *controlpanel :DelNetwork %s %s", username, network) +} diff --git a/irc/bncadmin/run.sh b/irc/bncadmin/run.sh new file mode 100644 index 0000000..9f8bd84 --- /dev/null +++ b/irc/bncadmin/run.sh @@ -0,0 +1,11 @@ +#!/bin/bash + +set -e +set -x + +box box.rb +docker push xena/bncadmin + +hyper rm -f bncadmin ||: +hyper pull xena/bncadmin +hyper run --name bncadmin -dit --size s1 --env-file .env xena/bncadmin diff --git a/irc/bncadmin/vendor-log b/irc/bncadmin/vendor-log new file mode 100644 index 0000000..90b30c8 --- /dev/null +++ b/irc/bncadmin/vendor-log @@ -0,0 +1,3 @@ +fd04337c94f98ab7c2ef34fbeb4e821284775095 github.com/belak/irc +4ed13390c0acd2ff4e371e64d8b97c8954138243 github.com/joho/godotenv +4ed13390c0acd2ff4e371e64d8b97c8954138243 github.com/joho/godotenv/autoload diff --git a/irc/bncadmin/vendor/github.com/belak/irc/client.go b/irc/bncadmin/vendor/github.com/belak/irc/client.go new file mode 100644 index 0000000..1eda9a8 --- /dev/null +++ b/irc/bncadmin/vendor/github.com/belak/irc/client.go @@ -0,0 +1,109 @@ +package irc + +import ( + "io" + "strings" +) + +// clientFilters are pre-processing which happens for certain message +// types. These were moved from below to keep the complexity of each +// component down. +var clientFilters = map[string]func(*Client, *Message){ + "001": func(c *Client, m *Message) { + c.currentNick = m.Params[0] + }, + "433": func(c *Client, m *Message) { + c.currentNick = c.currentNick + "_" + c.Writef("NICK :%s", c.currentNick) + }, + "437": func(c *Client, m *Message) { + c.currentNick = c.currentNick + "_" + c.Writef("NICK :%s", c.currentNick) + }, + "PING": func(c *Client, m *Message) { + reply := m.Copy() + reply.Command = "PONG" + c.WriteMessage(reply) + }, + "PRIVMSG": func(c *Client, m *Message) { + // Clean up CTCP stuff so everyone doesn't have to parse it + // manually. + lastArg := m.Trailing() + if len(lastArg) > 0 && lastArg[0] == '\x01' { + if i := strings.LastIndex(lastArg, "\x01"); i > -1 { + m.Command = "CTCP" + m.Params[len(m.Params)-1] = lastArg[1:i] + } + } + }, + "NICK": func(c *Client, m *Message) { + if m.Prefix.Name == c.currentNick && len(m.Params) > 0 { + c.currentNick = m.Params[0] + } + }, +} + +// ClientConfig is a structure used to configure a Client. +type ClientConfig struct { + // General connection information. + Nick string + Pass string + User string + Name string + + // Handler is used for message dispatching. + Handler Handler +} + +// Client is a wrapper around Conn which is designed to make common operations +// much simpler. +type Client struct { + *Conn + config ClientConfig + + // Internal state + currentNick string +} + +// NewClient creates a client given an io stream and a client config. +func NewClient(rwc io.ReadWriter, config ClientConfig) *Client { + return &Client{ + Conn: NewConn(rwc), + config: config, + } +} + +// Run starts the main loop for this IRC connection. Note that it may break in +// strange and unexpected ways if it is called again before the first connection +// exits. +func (c *Client) Run() error { + c.currentNick = c.config.Nick + + if c.config.Pass != "" { + c.Writef("PASS :%s", c.config.Pass) + } + + c.Writef("NICK :%s", c.config.Nick) + c.Writef("USER %s 0.0.0.0 0.0.0.0 :%s", c.config.User, c.config.Name) + + for { + m, err := c.ReadMessage() + if err != nil { + return err + } + + if f, ok := clientFilters[m.Command]; ok { + f(c, m) + } + + if c.config.Handler != nil { + c.config.Handler.Handle(c, m) + } + } +} + +// CurrentNick returns what the nick of the client is known to be at this point +// in time. +func (c *Client) CurrentNick() string { + return c.currentNick +} diff --git a/irc/bncadmin/vendor/github.com/belak/irc/conn.go b/irc/bncadmin/vendor/github.com/belak/irc/conn.go new file mode 100644 index 0000000..6c249a1 --- /dev/null +++ b/irc/bncadmin/vendor/github.com/belak/irc/conn.go @@ -0,0 +1,79 @@ +package irc + +import ( + "bufio" + "fmt" + "io" +) + +// Conn represents a simple IRC client. It embeds an irc.Reader and an +// irc.Writer. +type Conn struct { + *Reader + *Writer +} + +// NewConn creates a new Conn +func NewConn(rw io.ReadWriter) *Conn { + // Create the client + c := &Conn{ + NewReader(rw), + NewWriter(rw), + } + + return c +} + +// Writer is the outgoing side of a connection. +type Writer struct { + writer io.Writer +} + +// NewWriter creates an irc.Writer from an io.Writer. +func NewWriter(w io.Writer) *Writer { + return &Writer{w} +} + +// Write is a simple function which will write the given line to the +// underlying connection. +func (w *Writer) Write(line string) error { + _, err := w.writer.Write([]byte(line + "\r\n")) + return err +} + +// Writef is a wrapper around the connection's Write method and +// fmt.Sprintf. Simply use it to send a message as you would normally +// use fmt.Printf. +func (w *Writer) Writef(format string, args ...interface{}) error { + return w.Write(fmt.Sprintf(format, args...)) +} + +// WriteMessage writes the given message to the stream +func (w *Writer) WriteMessage(m *Message) error { + return w.Write(m.String()) +} + +// Reader is the incoming side of a connection. The data will be +// buffered, so do not re-use the io.Reader used to create the +// irc.Reader. +type Reader struct { + reader *bufio.Reader +} + +// NewReader creates an irc.Reader from an io.Reader. +func NewReader(r io.Reader) *Reader { + return &Reader{ + bufio.NewReader(r), + } +} + +// ReadMessage returns the next message from the stream or an error. +func (r *Reader) ReadMessage() (*Message, error) { + line, err := r.reader.ReadString('\n') + if err != nil { + return nil, err + } + + // Parse the message from our line + return ParseMessage(line) +} diff --git a/irc/bncadmin/vendor/github.com/belak/irc/handler.go b/irc/bncadmin/vendor/github.com/belak/irc/handler.go new file mode 100644 index 0000000..6a9fca7 --- /dev/null +++ b/irc/bncadmin/vendor/github.com/belak/irc/handler.go @@ -0,0 +1,16 @@ +package irc + +// Handler is a simple interface meant for dispatching a message from +// a Client connection. +type Handler interface { + Handle(*Client, *Message) +} + +// HandlerFunc is a simple wrapper around a function which allows it +// to be used as a Handler. +type HandlerFunc func(*Client, *Message) + +// Handle calls f(c, m) +func (f HandlerFunc) Handle(c *Client, m *Message) { + f(c, m) +} diff --git a/irc/bncadmin/vendor/github.com/belak/irc/parser.go b/irc/bncadmin/vendor/github.com/belak/irc/parser.go new file mode 100644 index 0000000..4204630 --- /dev/null +++ b/irc/bncadmin/vendor/github.com/belak/irc/parser.go @@ -0,0 +1,401 @@ +package irc + +import ( + "bytes" + "errors" + "strings" +) + +var tagDecodeSlashMap = map[rune]rune{ + ':': ';', + 's': ' ', + '\\': '\\', + 'r': '\r', + 'n': '\n', +} + +var tagEncodeMap = map[rune]string{ + ';': "\\:", + ' ': "\\s", + '\\': "\\\\", + '\r': "\\r", + '\n': "\\n", +} + +var ( + // ErrZeroLengthMessage is returned when parsing if the input is + // zero-length. + ErrZeroLengthMessage = errors.New("irc: Cannot parse zero-length message") + + // ErrMissingDataAfterPrefix is returned when parsing if there is + // no message data after the prefix. + ErrMissingDataAfterPrefix = errors.New("irc: No message data after prefix") + + // ErrMissingDataAfterTags is returned when parsing if there is no + // message data after the tags. + ErrMissingDataAfterTags = errors.New("irc: No message data after tags") + + // ErrMissingCommand is returned when parsing if there is no + // command in the parsed message. + ErrMissingCommand = errors.New("irc: Missing message command") +) + +// TagValue represents the value of a tag. +type TagValue string + +// ParseTagValue parses a TagValue from the connection. If you need to +// set a TagValue, you probably want to just set the string itself, so +// it will be encoded properly. +func ParseTagValue(v string) TagValue { + ret := &bytes.Buffer{} + + input := bytes.NewBufferString(v) + + for { + c, _, err := input.ReadRune() + if err != nil { + break + } + + if c == '\\' { + c2, _, err := input.ReadRune() + if err != nil { + ret.WriteRune(c) + break + } + + if replacement, ok := tagDecodeSlashMap[c2]; ok { + ret.WriteRune(replacement) + } else { + ret.WriteRune(c) + ret.WriteRune(c2) + } + } else { + ret.WriteRune(c) + } + } + + return TagValue(ret.String()) +} + +// Encode converts a TagValue to the format in the connection. +func (v TagValue) Encode() string { + ret := &bytes.Buffer{} + + for _, c := range v { + if replacement, ok := tagEncodeMap[c]; ok { + ret.WriteString(replacement) + } else { + ret.WriteRune(c) + } + } + + return ret.String() +} + +// Tags represents the IRCv3 message tags. +type Tags map[string]TagValue + +// ParseTags takes a tag string and parses it into a tag map. It will +// always return a tag map, even if there are no valid tags. +func ParseTags(line string) Tags { + ret := Tags{} + + tags := strings.Split(line, ";") + for _, tag := range tags { + parts := strings.SplitN(tag, "=", 2) + if len(parts) < 2 { + ret[parts[0]] = "" + continue + } + + ret[parts[0]] = ParseTagValue(parts[1]) + } + + return ret +} + +// GetTag is a convenience method to look up a tag in the map. +func (t Tags) GetTag(key string) (string, bool) { + ret, ok := t[key] + return string(ret), ok +} + +// Copy will create a new copy of all IRC tags attached to this +// message. +func (t Tags) Copy() Tags { + ret := Tags{} + + for k, v := range t { + ret[k] = v + } + + return ret +} + +// String ensures this is stringable +func (t Tags) String() string { + buf := &bytes.Buffer{} + + for k, v := range t { + buf.WriteByte(';') + buf.WriteString(k) + if v != "" { + buf.WriteByte('=') + buf.WriteString(v.Encode()) + } + } + + // We don't need the first byte because that's an extra ';' + // character. + buf.ReadByte() + + return buf.String() +} + +// Prefix represents the prefix of a message, generally the user who sent it +type Prefix struct { + // Name will contain the nick of who sent the message, the + // server who sent the message, or a blank string + Name string + + // User will either contain the user who sent the message or a blank string + User string + + // Host will either contain the host of who sent the message or a blank string + Host string +} + +// ParsePrefix takes an identity string and parses it into an +// identity struct. It will always return an Prefix struct and never +// nil. +func ParsePrefix(line string) *Prefix { + // Start by creating an Prefix with nothing but the host + id := &Prefix{ + Name: line, + } + + uh := strings.SplitN(id.Name, "@", 2) + if len(uh) == 2 { + id.Name, id.Host = uh[0], uh[1] + } + + nu := strings.SplitN(id.Name, "!", 2) + if len(nu) == 2 { + id.Name, id.User = nu[0], nu[1] + } + + return id +} + +// Copy will create a new copy of an Prefix +func (p *Prefix) Copy() *Prefix { + if p == nil { + return nil + } + + newPrefix := &Prefix{} + + *newPrefix = *p + + return newPrefix +} + +// String ensures this is stringable +func (p *Prefix) String() string { + buf := &bytes.Buffer{} + buf.WriteString(p.Name) + + if p.User != "" { + buf.WriteString("!") + buf.WriteString(p.User) + } + + if p.Host != "" { + buf.WriteString("@") + buf.WriteString(p.Host) + } + + return buf.String() +} + +// Message represents a line parsed from the server +type Message struct { + // Each message can have IRCv3 tags + Tags + + // Each message can have a Prefix + *Prefix + + // Command is which command is being called. + Command string + + // Params are all the arguments for the command. + Params []string +} + +// MustParseMessage calls ParseMessage and either returns the message +// or panics if an error is returned. +func MustParseMessage(line string) *Message { + m, err := ParseMessage(line) + if err != nil { + panic(err.Error()) + } + return m +} + +// ParseMessage takes a message string (usually a whole line) and +// parses it into a Message struct. This will return nil in the case +// of invalid messages. +func ParseMessage(line string) (*Message, error) { + // Trim the line and make sure we have data + line = strings.TrimSpace(line) + if len(line) == 0 { + return nil, ErrZeroLengthMessage + } + + c := &Message{ + Tags: Tags{}, + Prefix: &Prefix{}, + } + + if line[0] == '@' { + split := strings.SplitN(line, " ", 2) + if len(split) < 2 { + return nil, ErrMissingDataAfterTags + } + + c.Tags = ParseTags(split[0][1:]) + line = split[1] + } + + if line[0] == ':' { + split := strings.SplitN(line, " ", 2) + if len(split) < 2 { + return nil, ErrMissingDataAfterPrefix + } + + // Parse the identity, if there was one + c.Prefix = ParsePrefix(split[0][1:]) + line = split[1] + } + + // Split out the trailing then the rest of the args. Because + // we expect there to be at least one result as an arg (the + // command) we don't need to special case the trailing arg and + // can just attempt a split on " :" + split := strings.SplitN(line, " :", 2) + c.Params = strings.FieldsFunc(split[0], func(r rune) bool { + return r == ' ' + }) + + // If there are no args, we need to bail because we need at + // least the command. + if len(c.Params) == 0 { + return nil, ErrMissingCommand + } + + // If we had a trailing arg, append it to the other args + if len(split) == 2 { + c.Params = append(c.Params, split[1]) + } + + // Because of how it's parsed, the Command will show up as the + // first arg. + c.Command = c.Params[0] + c.Params = c.Params[1:] + + return c, nil +} + +// Trailing returns the last argument in the Message or an empty string +// if there are no args +func (m *Message) Trailing() string { + if len(m.Params) < 1 { + return "" + } + + return m.Params[len(m.Params)-1] +} + +// FromChannel is mostly for PRIVMSG messages (and similar derived +// messages) It will check if the message came from a channel or a +// person. This is not always accurate, because some servers allow for +// other characters at the start of a channel. +// +// TODO: Remove in favor of *Client.FromChannel(*Message) +func (m *Message) FromChannel() bool { + if len(m.Params) < 1 || len(m.Params[0]) < 1 { + return false + } + + switch m.Params[0][0] { + case '#', '&': + return true + default: + return false + } +} + +// Copy will create a new copy of an message +func (m *Message) Copy() *Message { + // Create a new message + newMessage := &Message{} + + // Copy stuff from the old message + *newMessage = *m + + // Copy any IRcv3 tags + newMessage.Tags = m.Tags.Copy() + + // Copy the Prefix + newMessage.Prefix = m.Prefix.Copy() + + // Copy the Params slice + newMessage.Params = append(make([]string, 0, len(m.Params)), m.Params...) + + return newMessage +} + +// String ensures this is stringable +func (m *Message) String() string { + buf := &bytes.Buffer{} + + // Write any IRCv3 tags if they exist in the message + if len(m.Tags) > 0 { + buf.WriteByte('@') + buf.WriteString(m.Tags.String()) + buf.WriteByte(' ') + } + + // Add the prefix if we have one + if m.Prefix != nil && m.Prefix.Name != "" { + buf.WriteByte(':') + buf.WriteString(m.Prefix.String()) + buf.WriteByte(' ') + } + + // Add the command since we know we'll always have one + buf.WriteString(m.Command) + + if len(m.Params) > 0 { + args := m.Params[:len(m.Params)-1] + trailing := m.Params[len(m.Params)-1] + + if len(args) > 0 { + buf.WriteByte(' ') + buf.WriteString(strings.Join(args, " ")) + } + + // If trailing is zero-length, contains a space or starts with + // a : we need to actually specify that it's trailing. + if len(trailing) == 0 || strings.ContainsRune(trailing, ' ') || trailing[0] == ':' { + buf.WriteString(" :") + } else { + buf.WriteString(" ") + } + buf.WriteString(trailing) + } + + return buf.String() +} diff --git a/irc/bncadmin/vendor/github.com/joho/godotenv/autoload/autoload.go b/irc/bncadmin/vendor/github.com/joho/godotenv/autoload/autoload.go new file mode 100644 index 0000000..fbcd2bd --- /dev/null +++ b/irc/bncadmin/vendor/github.com/joho/godotenv/autoload/autoload.go @@ -0,0 +1,15 @@ +package autoload + +/* + You can just read the .env file on import just by doing + + import _ "github.com/joho/godotenv/autoload" + + And bob's your mother's brother +*/ + +import "github.com/joho/godotenv" + +func init() { + godotenv.Load() +} diff --git a/irc/bncadmin/vendor/github.com/joho/godotenv/godotenv.go b/irc/bncadmin/vendor/github.com/joho/godotenv/godotenv.go new file mode 100644 index 0000000..94b2676 --- /dev/null +++ b/irc/bncadmin/vendor/github.com/joho/godotenv/godotenv.go @@ -0,0 +1,229 @@ +// Package godotenv is a go port of the ruby dotenv library (https://github.com/bkeepers/dotenv) +// +// Examples/readme can be found on the github page at https://github.com/joho/godotenv +// +// The TL;DR is that you make a .env file that looks something like +// +// SOME_ENV_VAR=somevalue +// +// and then in your go code you can call +// +// godotenv.Load() +// +// and all the env vars declared in .env will be avaiable through os.Getenv("SOME_ENV_VAR") +package godotenv + +import ( + "bufio" + "errors" + "os" + "os/exec" + "strings" +) + +// Load will read your env file(s) and load them into ENV for this process. +// +// Call this function as close as possible to the start of your program (ideally in main) +// +// If you call Load without any args it will default to loading .env in the current path +// +// You can otherwise tell it which files to load (there can be more than one) like +// +// godotenv.Load("fileone", "filetwo") +// +// It's important to note that it WILL NOT OVERRIDE an env variable that already exists - consider the .env file to set dev vars or sensible defaults +func Load(filenames ...string) (err error) { + filenames = filenamesOrDefault(filenames) + + for _, filename := range filenames { + err = loadFile(filename, false) + if err != nil { + return // return early on a spazout + } + } + return +} + +// Overload will read your env file(s) and load them into ENV for this process. +// +// Call this function as close as possible to the start of your program (ideally in main) +// +// If you call Overload without any args it will default to loading .env in the current path +// +// You can otherwise tell it which files to load (there can be more than one) like +// +// godotenv.Overload("fileone", "filetwo") +// +// It's important to note this WILL OVERRIDE an env variable that already exists - consider the .env file to forcefilly set all vars. +func Overload(filenames ...string) (err error) { + filenames = filenamesOrDefault(filenames) + + for _, filename := range filenames { + err = loadFile(filename, true) + if err != nil { + return // return early on a spazout + } + } + return +} + +// Read all env (with same file loading semantics as Load) but return values as +// a map rather than automatically writing values into env +func Read(filenames ...string) (envMap map[string]string, err error) { + filenames = filenamesOrDefault(filenames) + envMap = make(map[string]string) + + for _, filename := range filenames { + individualEnvMap, individualErr := readFile(filename) + + if individualErr != nil { + err = individualErr + return // return early on a spazout + } + + for key, value := range individualEnvMap { + envMap[key] = value + } + } + + return +} + +// Exec loads env vars from the specified filenames (empty map falls back to default) +// then executes the cmd specified. +// +// Simply hooks up os.Stdin/err/out to the command and calls Run() +// +// If you want more fine grained control over your command it's recommended +// that you use `Load()` or `Read()` and the `os/exec` package yourself. +func Exec(filenames []string, cmd string, cmdArgs []string) error { + Load(filenames...) + + command := exec.Command(cmd, cmdArgs...) + command.Stdin = os.Stdin + command.Stdout = os.Stdout + command.Stderr = os.Stderr + return command.Run() +} + +func filenamesOrDefault(filenames []string) []string { + if len(filenames) == 0 { + return []string{".env"} + } + return filenames +} + +func loadFile(filename string, overload bool) error { + envMap, err := readFile(filename) + if err != nil { + return err + } + + for key, value := range envMap { + if os.Getenv(key) == "" || overload { + os.Setenv(key, value) + } + } + + return nil +} + +func readFile(filename string) (envMap map[string]string, err error) { + file, err := os.Open(filename) + if err != nil { + return + } + defer file.Close() + + envMap = make(map[string]string) + + var lines []string + scanner := bufio.NewScanner(file) + for scanner.Scan() { + lines = append(lines, scanner.Text()) + } + + for _, fullLine := range lines { + if !isIgnoredLine(fullLine) { + key, value, err := parseLine(fullLine) + + if err == nil { + envMap[key] = value + } + } + } + return +} + +func parseLine(line string) (key string, value string, err error) { + if len(line) == 0 { + err = errors.New("zero length string") + return + } + + // ditch the comments (but keep quoted hashes) + if strings.Contains(line, "#") { + segmentsBetweenHashes := strings.Split(line, "#") + quotesAreOpen := false + var segmentsToKeep []string + for _, segment := range segmentsBetweenHashes { + if strings.Count(segment, "\"") == 1 || strings.Count(segment, "'") == 1 { + if quotesAreOpen { + quotesAreOpen = false + segmentsToKeep = append(segmentsToKeep, segment) + } else { + quotesAreOpen = true + } + } + + if len(segmentsToKeep) == 0 || quotesAreOpen { + segmentsToKeep = append(segmentsToKeep, segment) + } + } + + line = strings.Join(segmentsToKeep, "#") + } + + // now split key from value + splitString := strings.SplitN(line, "=", 2) + + if len(splitString) != 2 { + // try yaml mode! + splitString = strings.SplitN(line, ":", 2) + } + + if len(splitString) != 2 { + err = errors.New("Can't separate key from value") + return + } + + // Parse the key + key = splitString[0] + if strings.HasPrefix(key, "export") { + key = strings.TrimPrefix(key, "export") + } + key = strings.Trim(key, " ") + + // Parse the value + value = splitString[1] + // trim + value = strings.Trim(value, " ") + + // check if we've got quoted values + if strings.Count(value, "\"") == 2 || strings.Count(value, "'") == 2 { + // pull the quotes off the edges + value = strings.Trim(value, "\"'") + + // expand quotes + value = strings.Replace(value, "\\\"", "\"", -1) + // expand newlines + value = strings.Replace(value, "\\n", "\n", -1) + } + + return +} + +func isIgnoredLine(line string) bool { + trimmedLine := strings.Trim(line, " \n\t") + return len(trimmedLine) == 0 || strings.HasPrefix(trimmedLine, "#") +} |
