Random tweaks, codegen requests and responses

Author: christopher-dG

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -21,11 +21,11 @@ import (
     obs "github.com/christopher-dG/go-obs-websocket"
 )
 
-client := obs.NewClient("localhost", 4444, "")
+client := obs.Client{Host: "localhost", Port: 4444}
 if err := client.Connect(); err != nil {
     log.Fatal(err)
 }
-defer client.Close()
+defer client.Disconnect()
 
 // TODO
 ```

client.go

@@ -1,19 +1,96 @@
 package obsws
 
-import "github.com/gorilla/websocket"
+import (
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+	"strconv"
 
-// client is the interface to obs-websocket.
-type client struct {
+	"github.com/gorilla/websocket"
+	"github.com/pkg/errors"
+)
+
+// Client is the interface to obs-websocket.
+// Client{Host: "localhost", Port: 4444} will probably work if you haven't configured OBS.
+type Client struct {
 	Host     string // Host (probably "localhost").
 	Port     int    // Port (OBS default is 4444).
 	Password string // Password (OBS default is "").
 	conn     *websocket.Conn
 	id       int
 }
 
-// NewClient creates a new client. If you haven't configured obs-websocket at
-// all, then host should be "localhost", port should be 4444, and password
-// should be "".
-func NewClient(host string, port int, password string) *client {
-	return &client{Host: host, Port: port, Password: password}
+// Connect opens a WebSocket connection and authenticates if necessary.
+func (c *Client) Connect() error {
+	conn, err := connectWS(c.Host, c.Port)
+	if err != nil {
+		return err
+	}
+	c.conn = conn
+
+	reqGAR := GetAuthRequiredRequest{
+		MessageID:   c.getMessageID(),
+		RequestType: "GetAuthRequired",
+	}
+
+	if err = c.conn.WriteJSON(reqGAR); err != nil {
+		return errors.Wrap(err, "write Authenticate")
+	}
+
+	respGAR := &GetAuthRequiredResponse{}
+	if err = c.conn.ReadJSON(respGAR); err != nil {
+		return errors.Wrap(err, "read GetAuthRequired")
+	}
+
+	if !respGAR.AuthRequired {
+		logger.Info("no authentication required")
+		return nil
+	}
+
+	auth := getAuth(c.Password, respGAR.Salt, respGAR.Challenge)
+	logger.Debugf("auth: %s", auth)
+
+	reqA := AuthenticateRequest{
+		Auth: auth,
+		_request: _request{
+			MessageID:   c.getMessageID(),
+			RequestType: "Authenticate",
+		},
+	}
+	if err = c.conn.WriteJSON(reqA); err != nil {
+		return errors.Wrap(err, "write Authenticate")
+	}
+
+	logger.Info("logged in")
+	return nil
+}
+
+// Disconnect closes the WebSocket connection.
+func (c *Client) Disconnect() error {
+	return c.conn.Close()
+}
+
+func connectWS(host string, port int) (*websocket.Conn, error) {
+	url := fmt.Sprintf("ws://%s:%d", host, port)
+	logger.Infof("connecting to %s", url)
+	conn, _, err := websocket.DefaultDialer.Dial(url, nil)
+	if err != nil {
+		return nil, err
+	}
+	return conn, nil
+}
+
+func getAuth(password, salt, challenge string) string {
+	sha := sha256.Sum256([]byte(password + salt))
+	b64 := base64.StdEncoding.EncodeToString([]byte(sha[:]))
+
+	sha = sha256.Sum256([]byte(b64 + challenge))
+	b64 = base64.StdEncoding.EncodeToString([]byte(sha[:]))
+
+	return b64
+}
+
+func (c *Client) getMessageID() string {
+	c.id++
+	return strconv.Itoa(c.id)
 }

client_test.go

@@ -2,15 +2,10 @@ package obsws
 
 import "testing"
 
-func TestNewClient(t *testing.T) {
-	c := NewClient("localhost", 4444, "")
-	if c.Host != "localhost" {
-		t.Errorf("expected c.Host == 'localhost', got '%s'", c.Host)
-	}
-	if c.Port != 4444 {
-		t.Errorf("expected c.Port == 4444, got '%d'", c.Port)
-	}
-	if c.Password != "" {
-		t.Errorf("expected c.Password == '', got '%s'", c.Password)
+func TestGetAuth(t *testing.T) {
+	expected := "zTM5ki6L2vVvBQiTG9ckH1Lh64AbnCf6XZ226UmnkIA="
+	observed := getAuth("password", "salt", "challenge")
+	if observed != expected {
+		t.Errorf("expected auth == '%s', got '%s'", expected, observed)
 	}
 }

codegen/protocol.py

@@ -8,8 +8,10 @@
 package = "obsws"
 
 type_map = {
+    "bool": "bool",
     "boolean": "bool",
     "int": "int",
+    "float": "float64",
     "double": "float64",
     "string": "string",
     "array": "[]string",
@@ -36,36 +38,37 @@ def optional_type(s: str) -> Tuple[str, bool]:
 
 
 def process_json(d: Dict):
-    process_events(d["events"])
-    process_requests(d["requests"])
+    gen_events(d["events"])
+    gen_requests(d["requests"])
 
 
-def process_events(events: Dict):
+def gen_events(events: Dict):
+    """Generate all events."""
     for category, data in events.items():
-        process_events_category(category, data)
+        gen_events_category(category, data)
 
 
-def process_events_category(category: str, data: Dict):
-    events = "\n\n".join(generate_event(event) for event in data)
+def gen_events_category(category: str, data: Dict):
+    """Generate all events in one category."""
+    events = "\n\n".join(gen_event(event) for event in data)
     with open(go_filename("events", category), "w") as f:
         f.write(f"""\
         package {package}
 
-        // This code is automatically generated.
-        // See: https://github.com/christopher-dG/go-obs-websocket/blob/master/codegen/protocol.py
-
-        // https://github.com/Palakis/obs-websocket/blob/master/docs/generated/protocol.md#{category}
+        // This file is automatically generated.
+        // https://github.com/christopher-dG/go-obs-websocket/blob/master/codegen/protocol.py
 
         {events}
         """)
 
 
-def generate_event(data: Dict) -> str:
-    """Generate Go code with type definitions and interface functions."""
-    if "returns" in data:
+def gen_event(data: Dict) -> str:
+    """Write Go code with a type definition and interface functions."""
+    reserved = ["Type", "StreamTC", "RecTC"]
+    if data.get("returns"):
         struct = f"""\
         type {data['name']}Event struct {{
-            {go_variables(data['returns'])}
+            {go_variables(data['returns'], reserved)}
             _event
         }}\
         """
@@ -74,8 +77,10 @@ def generate_event(data: Dict) -> str:
 
     description = data["description"].replace("\n", " ")
     description = f"{data['name']}Event : {description}"
-    if "since" in data:
-        description += f" Since: {data['since'].capitalize()}"
+    if description and not description.endswith("."):
+        description += "."
+    if data.get("since"):
+        description += f" Since: {data['since'].capitalize()}."
 
     return f"""\
     // {description}
@@ -93,44 +98,135 @@ def generate_event(data: Dict) -> str:
     """
 
 
-def process_requests(requests: Dict):
-    pass
+def gen_requests(requests: Dict):
+    """Generate all requests and responses."""
+    for category, data in requests.items():
+        gen_requests_category(category, data)
+
+
+def gen_requests_category(category: str, data: Dict):
+    requests = "\n\n".join(gen_request(request) for request in data)
+    with open(go_filename("requests", category), "w") as f:
+        f.write(f"""\
+        package {package}
 
+        // This file is automatically generated.
+        // https://github.com/christopher-dG/go-obs-websocket/blob/master/codegen/protocol.py
 
-def go_variables(vars: List) -> str:
+        {requests}
+        """)
+
+
+def gen_request(data: Dict) -> str:
+    """Write Go code with type definitions and interface functions."""
+    reserved = ["ID", "Type"]
+    if data.get("params"):
+        struct = f"""\
+        type {data['name']}Request struct {{
+            {go_variables(data['params'], reserved)}
+            _request
+        }}
+        """
+    else:
+        struct = f"type {data['name']}Request _request"
+
+    description = data["description"].replace("\n", " ")
+    description = f"{data['name']}Request : {description}"
+    if description and not description.endswith("."):
+        description += "."
+    if data.get("since"):
+        description += f" Since: {data['since'].capitalize()}."
+
+    request = f"""\
+    // {description}
+    // https://github.com/Palakis/obs-websocket/blob/master/docs/generated/protocol.md#{data['heading']['text'].lower()}
+    {struct}
+
+    // ID returns the request's message ID.
+    func (r {data['name']}Request) ID() string {{ return r.MessageID }}
+
+    // Type returns the request's message type.
+    func (r {data['name']}Request) Type() string {{ return r.RequestType }}
     """
-    Convert a list of variable definition into Go code to be put
+
+    if data.get("returns"):
+        reserved = ["ID", "Stat", "Err"]
+        struct = f"""\
+        type {data['name']}Response struct {{
+            {go_variables(data['returns'], reserved)}
+            _response
+        }}
+        """
+    else:
+        struct = f"type {data['name']}Response _response"
+
+    description = f"{data['name']}Response : Response for {data['name']}Request."
+    if data.get("since"):
+        description += f" Since: {data['since'].capitalize()}."
+
+    response = f"""\
+    // {description}
+    // https://github.com/Palakis/obs-websocket/blob/master/docs/generated/protocol.md#{data['heading']['text'].lower()}
+    {struct}
+
+    // ID returns the response's message ID.
+    func (r {data['name']}Response) ID() string {{ return r.MessageID }}
+
+    // Stat returns the response's status.
+    func (r {data['name']}Response) Stat() string {{ return r.Status }}
+
+    // Err returns the response's error.
+    func (r {data['name']}Response) Err() string {{ return r.Error }}
+    """
+
+    return f"{request}\n\n{response}"
+
+
+def go_variables(names: List, reserved: List) -> str:
+    """
+    Convert a list of variable names into Go code to be put
     inside a struct definition.
     """
-    lines = []
-    for v in vars:
-        line = go_name(v["name"])
+    lines, varnames = [], []
+    for v in names:
         typename, optional = optional_type(v["type"])
-        line += f" {type_map[typename.lower()]} // {v['description']}"
+        varname = go_var(v["name"])
+        description = v["description"].replace("\n", " ")
+        if description and not description.endswith("."):
+            description += "."
+        tag = '`json:"%s"`' % v['name']
+        line = f"{go_var(v['name'])} {type_map[typename.lower()]} {tag} // {description}"
         if optional:
-            line += " Optional."
+            line += " Optional." if description else "Optional."
+        if varname in reserved:
+            line += " TODO: Reserved name."
+        if varname in varnames:
+            line += " TODO: Duplicate name."
+        else:
+            varnames.append(varname)
         if typename.lower() in unknown_types:
-            line += f" TODO: Unknown type ({typename})."
+            line += f" TODO: Unknown type ({v['type']})."
         lines.append(line)
     return "\n".join(lines)
 
 
-def go_name(s: str) -> str:
-    """
-    Convert a variable name in the input file to a Go variable name.
-    Note: This makes lots of assumptions about the input,
-    i.e. nothing ends with a separator.
-    """
-    s = s.capitalize()
+def go_var(s: str) -> str:
+    """Convert a variable name in the input file to a Go variable name."""
+    s = f"{s[0].upper()}{s[1:]}"
     for sep in ["-", "_", ".*.", "[].", "."]:
         while sep in s:
-            i = s.find(sep)
             _len = len(sep)
+            if s.endswith(sep):
+                s = s[:-_len]
+                continue
+            i = s.find(sep)
             s = f"{s[:i]}{s[i+_len].upper()}{s[i+_len+1:]}"
-    return s
+
+    return s.replace("Id", "ID")  # Yuck.
 
 
 def go_filename(category, section):
+    """Generate a Go filename from a category and section."""
     return f"{category}_{section.replace(' ', '_')}.go"