go-tg-screenshot-bot

1
package dbus
2

3
import (
4
	"bufio"
5
	"bytes"
6
	"errors"
7
	"io"
8
	"os"
9
	"strconv"
10
)
11

12
// AuthStatus represents the Status of an authentication mechanism.
13
type AuthStatus byte
14

15
const (
16
	// AuthOk signals that authentication is finished; the next command
17
	// from the server should be an OK.
18
	AuthOk AuthStatus = iota
19

20
	// AuthContinue signals that additional data is needed; the next command
21
	// from the server should be a DATA.
22
	AuthContinue
23

24
	// AuthError signals an error; the server sent invalid data or some
25
	// other unexpected thing happened and the current authentication
26
	// process should be aborted.
27
	AuthError
28
)
29

30
type authState byte
31

32
const (
33
	waitingForData authState = iota
34
	waitingForOk
35
	waitingForReject
36
)
37

38
// Auth defines the behaviour of an authentication mechanism.
39
type Auth interface {
40
	// Return the name of the mechanism, the argument to the first AUTH command
41
	// and the next status.
42
	FirstData() (name, resp []byte, status AuthStatus)
43

44
	// Process the given DATA command, and return the argument to the DATA
45
	// command and the next status. If len(resp) == 0, no DATA command is sent.
46
	HandleData(data []byte) (resp []byte, status AuthStatus)
47
}
48

49
// Auth authenticates the connection, trying the given list of authentication
50
// mechanisms (in that order). If nil is passed, the EXTERNAL and
51
// DBUS_COOKIE_SHA1 mechanisms are tried for the current user. For private
52
// connections, this method must be called before sending any messages to the
53
// bus. Auth must not be called on shared connections.
54
func (conn *Conn) Auth(methods []Auth) error {
55
	if methods == nil {
56
		uid := strconv.Itoa(os.Geteuid())
57
		methods = []Auth{AuthExternal(uid), AuthCookieSha1(uid, getHomeDir())}
58
	}
59
	in := bufio.NewReader(conn.transport)
60
	err := conn.transport.SendNullByte()
61
	if err != nil {
62
		return err
63
	}
64
	err = authWriteLine(conn.transport, []byte("AUTH"))
65
	if err != nil {
66
		return err
67
	}
68
	s, err := authReadLine(in)
69
	if err != nil {
70
		return err
71
	}
72
	if len(s) < 2 || !bytes.Equal(s[0], []byte("REJECTED")) {
73
		return errors.New("dbus: authentication protocol error")
74
	}
75
	s = s[1:]
76
	for _, v := range s {
77
		for _, m := range methods {
78
			if name, _, status := m.FirstData(); bytes.Equal(v, name) {
79
				var ok bool
80
				err = authWriteLine(conn.transport, []byte("AUTH"), v)
81
				if err != nil {
82
					return err
83
				}
84
				switch status {
85
				case AuthOk:
86
					err, ok = conn.tryAuth(m, waitingForOk, in)
87
				case AuthContinue:
88
					err, ok = conn.tryAuth(m, waitingForData, in)
89
				default:
90
					panic("dbus: invalid authentication status")
91
				}
92
				if err != nil {
93
					return err
94
				}
95
				if ok {
96
					if conn.transport.SupportsUnixFDs() {
97
						err = authWriteLine(conn, []byte("NEGOTIATE_UNIX_FD"))
98
						if err != nil {
99
							return err
100
						}
101
						line, err := authReadLine(in)
102
						if err != nil {
103
							return err
104
						}
105
						switch {
106
						case bytes.Equal(line[0], []byte("AGREE_UNIX_FD")):
107
							conn.EnableUnixFDs()
108
							conn.unixFD = true
109
						case bytes.Equal(line[0], []byte("ERROR")):
110
						default:
111
							return errors.New("dbus: authentication protocol error")
112
						}
113
					}
114
					err = authWriteLine(conn.transport, []byte("BEGIN"))
115
					if err != nil {
116
						return err
117
					}
118
					go conn.inWorker()
119
					return nil
120
				}
121
			}
122
		}
123
	}
124
	return errors.New("dbus: authentication failed")
125
}
126

127
// tryAuth tries to authenticate with m as the mechanism, using state as the
128
// initial authState and in for reading input. It returns (nil, true) on
129
// success, (nil, false) on a REJECTED and (someErr, false) if some other
130
// error occurred.
131
func (conn *Conn) tryAuth(m Auth, state authState, in *bufio.Reader) (error, bool) {
132
	for {
133
		s, err := authReadLine(in)
134
		if err != nil {
135
			return err, false
136
		}
137
		switch {
138
		case state == waitingForData && string(s[0]) == "DATA":
139
			if len(s) != 2 {
140
				err = authWriteLine(conn.transport, []byte("ERROR"))
141
				if err != nil {
142
					return err, false
143
				}
144
				continue
145
			}
146
			data, status := m.HandleData(s[1])
147
			switch status {
148
			case AuthOk, AuthContinue:
149
				if len(data) != 0 {
150
					err = authWriteLine(conn.transport, []byte("DATA"), data)
151
					if err != nil {
152
						return err, false
153
					}
154
				}
155
				if status == AuthOk {
156
					state = waitingForOk
157
				}
158
			case AuthError:
159
				err = authWriteLine(conn.transport, []byte("ERROR"))
160
				if err != nil {
161
					return err, false
162
				}
163
			}
164
		case state == waitingForData && string(s[0]) == "REJECTED":
165
			return nil, false
166
		case state == waitingForData && string(s[0]) == "ERROR":
167
			err = authWriteLine(conn.transport, []byte("CANCEL"))
168
			if err != nil {
169
				return err, false
170
			}
171
			state = waitingForReject
172
		case state == waitingForData && string(s[0]) == "OK":
173
			if len(s) != 2 {
174
				err = authWriteLine(conn.transport, []byte("CANCEL"))
175
				if err != nil {
176
					return err, false
177
				}
178
				state = waitingForReject
179
			} else {
180
				conn.uuid = string(s[1])
181
				return nil, true
182
			}
183
		case state == waitingForData:
184
			err = authWriteLine(conn.transport, []byte("ERROR"))
185
			if err != nil {
186
				return err, false
187
			}
188
		case state == waitingForOk && string(s[0]) == "OK":
189
			if len(s) != 2 {
190
				err = authWriteLine(conn.transport, []byte("CANCEL"))
191
				if err != nil {
192
					return err, false
193
				}
194
				state = waitingForReject
195
			} else {
196
				conn.uuid = string(s[1])
197
				return nil, true
198
			}
199
		case state == waitingForOk && string(s[0]) == "DATA":
200
			err = authWriteLine(conn.transport, []byte("DATA"))
201
			if err != nil {
202
				return err, false
203
			}
204
		case state == waitingForOk && string(s[0]) == "REJECTED":
205
			return nil, false
206
		case state == waitingForOk && string(s[0]) == "ERROR":
207
			err = authWriteLine(conn.transport, []byte("CANCEL"))
208
			if err != nil {
209
				return err, false
210
			}
211
			state = waitingForReject
212
		case state == waitingForOk:
213
			err = authWriteLine(conn.transport, []byte("ERROR"))
214
			if err != nil {
215
				return err, false
216
			}
217
		case state == waitingForReject && string(s[0]) == "REJECTED":
218
			return nil, false
219
		case state == waitingForReject:
220
			return errors.New("dbus: authentication protocol error"), false
221
		default:
222
			panic("dbus: invalid auth state")
223
		}
224
	}
225
}
226

227
// authReadLine reads a line and separates it into its fields.
228
func authReadLine(in *bufio.Reader) ([][]byte, error) {
229
	data, err := in.ReadBytes('\n')
230
	if err != nil {
231
		return nil, err
232
	}
233
	data = bytes.TrimSuffix(data, []byte("\r\n"))
234
	return bytes.Split(data, []byte{' '}), nil
235
}
236

237
// authWriteLine writes the given line in the authentication protocol format
238
// (elements of data separated by a " " and terminated by "\r\n").
239
func authWriteLine(out io.Writer, data ...[]byte) error {
240
	buf := make([]byte, 0)
241
	for i, v := range data {
242
		buf = append(buf, v...)
243
		if i != len(data)-1 {
244
			buf = append(buf, ' ')
245
		}
246
	}
247
	buf = append(buf, '\r')
248
	buf = append(buf, '\n')
249
	n, err := out.Write(buf)
250
	if err != nil {
251
		return err
252
	}
253
	if n != len(buf) {
254
		return io.ErrUnexpectedEOF
255
	}
256
	return nil
257
}
258