1package policy
 2
 3import (
 4	"errors"
 5	"net"
 6
 7	"go.starlark.net/starlark"
 8	"go.starlark.net/starlarkstruct"
 9)
10
11type Conn struct {
12	RemoteAddr starlark.String
13	ClientIP   starlark.String
14	Meta       *starlarkstruct.Struct
15}
16
17func (c Conn) String() string       { return "<Conn>" }
18func (c Conn) Type() string         { return "Conn" }
19func (c Conn) Freeze()              { c.Meta.Freeze() }
20func (c Conn) Truth() starlark.Bool { return true }
21func (c Conn) Hash() (uint32, error) {
22	return 0, errors.New("unhashable")
23}
24
25func (c Conn) Attr(name string) (starlark.Value, error) {
26	switch name {
27	case "remote_addr":
28		return c.RemoteAddr, nil
29	case "client_ip":
30		return c.ClientIP, nil
31	case "meta":
32		return c.Meta, nil
33	}
34	return nil, nil
35}
36
37func (c Conn) AttrNames() []string {
38	return []string{"remote_addr", "client_ip", "meta"}
39}
40
41func ConnToStarlark(conn net.Conn) Conn {
42	remoteAddr := conn.RemoteAddr().String()
43	var clientIP string
44	if tcpAddr, ok := conn.RemoteAddr().(*net.TCPAddr); ok {
45		clientIP = tcpAddr.IP.String()
46	} else {
47		clientIP = remoteAddr
48	}
49
50	return Conn{
51		RemoteAddr: starlark.String(remoteAddr),
52		ClientIP:   starlark.String(clientIP),
53		Meta:       starlarkstruct.FromStringDict(starlark.None, starlark.StringDict{}),
54	}
55}