tracer/tracer.py

#!/usr/bin/env python3

from random import randint, choice
from tracery.modifiers import base_english
import configparser
import glob
import irctokens
import json
import os
import random
import re
import socket
import subprocess
import sys
import time
import traceback
import tracery

DB = {}
config = configparser.ConfigParser(
    converters={"list": lambda x: [i.strip() for i in x.split(",")]}
)
config.read("config.ini")
bot = config["irc"]

# read account info if it exists
if os.path.isfile("account.ini"):
    account = configparser.ConfigParser()
    account.read("account.ini")
    account = account["nickserv"]


def grammar(rules):
    try:
        res = tracery.Grammar(rules)
        res.add_modifiers(base_english)
        return res
    except Exception as e:
        print(e)


def load_rules(path):
    try:
        with open(path) as f:
            return json.loads(f.read())
    except Exception as e:
        print(e)


def populate():
    global DB
    DB = {}
    for p in glob.glob("/home/*/.tracery/*"):
        name, ext = os.path.splitext(p)
        name = os.path.basename(name)
        if name.startswith(".") or ext not in (".json", ""):
            continue
        if p in DB:
            DB[name].append(grammar(load_rules(p)))
        else:
            DB[name] = [grammar(load_rules(p))]


populate()


def generate(rule):
    populate()
    if rule in DB:
        g = random.choice(DB[rule])
        return g.flatten("#origin#")


def listify(col):
    if type(col) == type([]):
        return col
    else:
        return [col]


def shuffle(col):
    a = random.choice(list(col))
    b = random.choice(list(col))
    if "origin" in [a, b]:
        return col
    col[a], col[b] = col[b], col[a]
    return col


def fuse(argv):
    populate()
    raw = {}
    for gk in argv:
        if gk in DB:
            g = random.choice(DB[gk]).raw
            for k in g:
                if k in raw:
                    raw[k] = listify(raw[k]) + listify(g[k])
                else:
                    raw[k] = g[k]
    for i in range(20):
        raw = shuffle(raw)
    return grammar(raw).flatten("#origin#")


def _send(line):
    print(f"> {line.format()}")
    e.push(line)
    while e.pending():
        e.pop(s.send(e.pending()))


def send(chan, msg):
    _send(irctokens.build("PRIVMSG", [chan, msg]))


def think(line):
    chan = line.params.pop(0)
    words = line.params[0].split(" ")

    if len(words) > 0 and line.hostmask.nickname != bot["nick"]:
        if words[0] == "!!list":
            res = ""
            for k in DB:
                res += k + " "
            send(chan, res[:475])
        elif words[0] == "!!fuse":
            if "|" in words:
                res = fuse(words[1 : words.index("|")])
                if res:
                    send(chan, " ".join(words[words.index("|") + 1 :]) + " " + res)
            else:
                res = fuse(words[1:])
                if res:
                    send(chan, res[0:475])
        elif words[0] == "!!source":
            send(chan, "https://tildegit.org/ben/tracer")
        elif words[0] == "!botlist" or words[0] == "!!help":
            send(
                chan,
                "helo i'm a tracery bot that makes cool things from tracery grammars in your ~/.tracery. see http://tracery.io for more info",
            )
        elif words[0][0:2] == "!!":
            print(words)
            res = generate(words[0][2:])
            if res:
                if len(words) >= 3:
                    if words[1] == "|":
                        send(chan, " ".join(words[2:]) + " " + res)
                    else:
                        send(chan, res)
                else:
                    send(chan, res)


if __name__ == "__main__":
    d = irctokens.StatefulDecoder()
    e = irctokens.StatefulEncoder()
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((bot["server"], int(bot["port"])))

    _send(irctokens.build("USER", [bot["nick"], "0", "*", "tracery bot"]))
    _send(irctokens.build("NICK", [bot["nick"]]))

    while True:
        lines = d.push(s.recv(1024))

        if lines == None:
            print("! disconnected")
            break

        for line in lines:
            print(f"< {line.format()}")

            if line.command == "PING":
                _send(irctokens.build("PONG", [line.params[0]]))

            elif line.command == "001":
                _send(irctokens.build("MODE", [bot["nick"], "+B"]))
                if account is not None:
                    _send(
                        irctokens.build(
                            "SQUERY",
                            [
                                "NickServ",
                                "IDENTIFY",
                                account["username"],
                                account["password"],
                            ],
                        )
                    )

                _send(irctokens.build("JOIN", [",".join(bot.getlist("channels"))]))

            elif line.command == "INVITE":
                _send(irctokens.build("JOIN", [line.params[1]]))

            elif line.command == "PRIVMSG":
                try:
                    think(line)
                except Exception as e:
                    print("ERROR", line)
                    print(e)
                    traceback.print_exc()