pnut-matrix/pnut-matrix.py
2021-03-20 13:26:33 +00:00

408 lines
13 KiB
Python

import websocket
import threading
import time
import logging
import yaml
import json
import pnutpy
import requests
import magic
import argparse
import os
from matrix_client.api import MatrixHttpApi
from matrix_client.api import MatrixError, MatrixRequestError
from models import Avatars, Rooms, Events, DirectRooms, Users
from database import db_session, init_db
from sqlalchemy import and_
from appservice import app
logger = logging.getLogger()
_shutdown = threading.Event()
_reconnect = threading.Event()
def new_message(msg, meta):
logger.debug("channel: " + msg.channel_id)
logger.debug("username: " + msg.user.username)
if 'name' in msg.user:
logger.debug("name: " + msg.user.name)
logger.debug("text: " + msg.content.text)
# ignore messages posted by the bridge
if msg.user.username == config['MATRIX_PNUT_USER']:
return
if msg.source.id == config['PNUTCLIENT_ID']:
return
if meta['channel_type'] == 'io.pnut.core.chat':
room = Rooms.query.filter(Rooms.pnut_chan == msg.channel_id).one_or_none()
elif meta['channel_type'] == 'io.pnut.core.pm':
room = DirectRooms.query.filter(DirectRooms.pnut_chan == msg.channel_id).one_or_none()
if room is None:
# Do do an invite from the bridge user?
logger.debug('new invite?')
# create room and included matrix recpient
# subscribed_user_ids from meta
logger.debug(meta['subscribed_user_ids'])
pnut_user = matrix_id_from_pnut(msg.user.username)
profile = get_matrix_profile(pnut_user)
if not profile:
new_matrix_user(msg.user.username)
invitees=[]
for pm_user in meta['subscribed_user_ids']:
user = Users.query.filter(Users.pnut_user_id == pm_user).one_or_none()
if user is not None:
invitees.append(user.matrix_id)
if len(invitees) > 0:
room = new_room(pnut_user, invitees, msg.channel_id)
else:
room = None
logger.debug(room)
if room is None:
logger.debug('-not_mapped-')
return
matrix_id = matrix_id_from_pnut(msg.user.username)
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=matrix_id)
profile = get_matrix_profile(matrix_id)
if not profile:
new_matrix_user(msg.user.username)
logger.debug('-new_user-')
profile = {'displayname': None}
if profile['displayname'] != matrix_display_from_pnut(msg.user):
set_matrix_display(msg.user)
logger.debug('-set_display-')
avatar = Avatars.query.filter(Avatars.pnut_user == msg.user.username).one_or_none()
if avatar is None or avatar.avatar != msg.user.content.avatar_image.url:
set_matrix_avatar(msg.user)
logger.debug('-set_avatar-')
# members = matrix_api.get_room_members(room.room_id)
# logger.debug(members)
# join_room(room.room_id, config['MATRIX_AS_ID'])
# TODO: sort out room invite and join logic
join_room(room.room_id, matrix_id)
if 'content' in msg:
text = msg.content.text + "\n"
ts = int(msg.created_at.strftime('%s')) * 1000
lnktext = ""
for link in msg.content.entities.links:
if 'title' in link:
lnktext += link.title + "\n"
if 'url' in link:
lnktext += link.url + "\n"
if len(lnktext) > 0:
text += "\n" + lnktext
r = matrix_api.send_message(room.room_id, text, timestamp=ts)
event = Events(
event_id=r['event_id'],
room_id=room.room_id,
pnut_msg_id=msg.id,
pnut_user_id=msg.user.id,
pnut_chan_id=msg.channel_id,
deleted=False)
db_session.add(event)
db_session.commit()
if 'raw' in msg:
logger.debug('-handle media uploads-')
new_media(room.room_id, msg)
def new_media(room_id, msg):
matrix_id = matrix_id_from_pnut(msg.user.username)
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=matrix_id)
ts = int(msg.created_at.strftime('%s')) * 1000
if 'io.pnut.core.oembed' in msg.raw:
for oembed in msg.raw['io.pnut.core.oembed']:
info = {}
if oembed.type == 'photo':
msgtype = 'm.image'
dl_url = oembed.url
info['h'] = oembed.height
info['w'] = oembed.width
elif oembed.type == 'audio':
logger.debug("* recieved audio attachment")
continue
elif oembed.type == 'video':
logger.debug("* recieved video attachment")
continue
elif oembed.type == 'html5video':
logger.debug("* recieved html5 video attachment")
continue
elif oembed.type == 'rich':
logger.debug("* recieved video attachment")
continue
else:
logger.debug("* recieved unknown attachment")
continue
dl = requests.get(dl_url, stream=True)
dl.raise_for_status()
with magic.Magic(flags=magic.MAGIC_MIME_TYPE) as m:
info['mimetype'] = m.id_buffer(dl.content)
info['size'] = len(dl.content)
ul = matrix_api.media_upload(dl.content, info['mimetype'])
if 'title' in oembed:
title = oembed.title
else:
title = ""
r = matrix_api.send_content(room_id, ul['content_uri'], title, msgtype, extra_information=info, timestamp=ts)
event = Events(
event_id=r['event_id'],
room_id=room_id,
pnut_msg_id=msg.id,
pnut_user_id=msg.user.id,
pnut_chan_id=msg.channel_id,
deleted=False)
db_session.add(event)
db_session.commit()
def delete_message(msg):
matrix_id = matrix_id_from_pnut(msg.user.username)
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=matrix_id)
events = Events.query.filter(and_(Events.pnut_msg_id == msg.id, Events.deleted == False)).all()
for event in events:
matrix_api.redact_event(event.room_id, event.event_id)
event.deleted = True
db_session.commit()
def matrix_id_from_pnut(username):
return "@" + config['MATRIX_PNUT_PREFIX'] + username + ":" + config['MATRIX_DOMAIN']
def matrix_display_from_pnut(user):
if 'name' in user:
display = user.name + " <@" + user.username + "> (pnut)"
else:
display = "@" + user.username
return display
# return user.username + " (pnut)"
def get_matrix_profile(matrix_id):
url = matrix_url + '/profile/' + matrix_id
r = requests.get(url)
if r.status_code == 200:
return r.json()
else:
return None
def set_matrix_display(user):
matrix_id = matrix_id_from_pnut(user.username)
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=matrix_id)
matrix_api.set_display_name(matrix_id, matrix_display_from_pnut(user))
def set_matrix_avatar(user):
matrix_id = matrix_id_from_pnut(user.username)
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=matrix_id)
dl = requests.get(user.content.avatar_image.url, stream=True)
dl.raise_for_status()
with magic.Magic(flags=magic.MAGIC_MIME_TYPE) as m:
mtype = m.id_buffer(dl.content)
ul = matrix_api.media_upload(dl.content, mtype)
try:
matrix_api.set_avatar_url(matrix_id, ul['content_uri'])
avatar = Avatars.query.filter(Avatars.pnut_user == user.username).one_or_none()
if avatar is None:
avatar = Avatars(pnut_user=user.username, avatar=user.content.avatar_image.url)
db_session.add(avatar)
else:
avatar.avatar = user.content.avatar_image.url
db_session.commit()
except MatrixRequestError:
logger.exception('failed to set user avatar')
def new_matrix_user(username):
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'])
data = {
'type': 'm.login.application_service',
'user': config['MATRIX_PNUT_PREFIX'] + username
}
matrix_api.register(content=data)
def join_room(room_id, matrix_id):
matrix_api_as = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'])
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=matrix_id)
try:
matrix_api.join_room(room_id)
except MatrixRequestError as e:
if e.code == 403:
matrix_api_as.invite_user(room_id, matrix_id)
matrix_api.join_room(room_id)
else:
logger.exception('failed to join room')
logger.debug('-room_join-')
def new_room(pnut_user, invitees, chan):
dr = None
matrix_api = MatrixHttpApi(config['MATRIX_HOST'],
token=config['MATRIX_AS_TOKEN'],
identity=pnut_user)
url = matrix_url + '/createRoom'
params = {"access_token": config['MATRIX_AS_TOKEN'], "user_id": pnut_user}
content = {"visibility": "private", "is_direct": True, "invite": invitees}
headers = {"Content-Type": "application/json"}
r = requests.post(url, headers=headers, params=params,
data=json.dumps(content))
response = r.json()
for bridge_user in invitees:
dr = DirectRooms(room_id=response['room_id'],
bridge_user=pnut_user, pnut_chan=chan)
logger.debug(dr)
db_session.add(dr)
db_session.commit()
return dr
def on_message(ws, message):
# logger.debug("on_message: " + message)
msg = json.loads(message)
logger.debug(msg['meta'])
if 'data' in msg:
if 'channel_type' in msg['meta']:
if msg['meta']['channel_type'] not in ['io.pnut.core.chat',
'io.pnut.core.pm']:
return
for d_item in msg['data']:
pmsg = pnutpy.models.Message.from_response_data(d_item)
if 'is_deleted' in msg['meta']:
if msg['meta']['is_deleted']:
logger.debug("message: delete")
delete_message(pmsg)
else:
new_message(pmsg, msg['meta'])
def on_error(ws, error):
logger.error("on_error: !!! ERROR !!!")
logger.error(error)
def on_close(ws):
logger.debug("on_close: ### CLOSED ###")
def on_open(ws):
def run(*args):
while not _shutdown.isSet() and not _reconnect.isSet():
time.sleep(3)
try:
ws.send(".")
except websocket._exceptions.WebSocketConnectionClosedException:
logger.debug('websocket closed exception caught...')
_reconnect.set()
time.sleep(1)
logger.debug("*** terminate thread ***")
t = threading.Thread(target=run)
t.start()
def wsthreader(threadfunc):
def wrapper():
while not _shutdown.isSet():
_reconnect.clear()
logger.debug('threadfunc start...')
running = threadfunc()
logger.debug('threadfunc end...')
if running:
time.sleep(5)
else:
_shutdown.set()
logger.debug('*** thread stopped ***')
return wrapper
if __name__ == '__main__':
a_parser = argparse.ArgumentParser()
a_parser.add_argument(
'-d', action='store_true', dest='debug',
help="debug logging"
)
# TODO: solve the database.py problem and enable this
# a_parser.add_argument(
# '-c', '--config', default="config.yaml",
# help="configuration file"
# )
args = a_parser.parse_args()
if args.debug:
# websocket.enableTrace(True)
logging.basicConfig(level=logging.DEBUG)
else:
logging.basicConfig(level=logging.INFO)
configyaml = os.environ.get("CONFIG_FILE")
with open(configyaml, "rb") as config_file:
config = yaml.load(config_file, Loader=yaml.SafeLoader)
ws_url = 'wss://stream.pnut.io/v1/app?access_token='
ws_url += config['PNUT_APPTOKEN'] + '&key=' + config['PNUT_APPKEY']
ws_url += '&include_raw=1'
matrix_url = config['MATRIX_HOST'] + '/_matrix/client/r0'
# setup the database connection
init_db()
# setup the websocket connection
ws = websocket.WebSocketApp(ws_url, on_message=on_message,
on_error=on_error, on_close=on_close)
ws.on_open = on_open
wst = threading.Thread(target=wsthreader(ws.run_forever))
wst.start()
# setup the matrix app service
if config['MATRIX_ADMIN_ROOM']:
logger.debug("- sould join admin room -")
join_room(config['MATRIX_ADMIN_ROOM'], config['MATRIX_AS_ID'])
app.config.update(config)
app.run(host=config['LISTEN_HOST'], port=config['LISTEN_PORT'])
logger.info('!! shutdown initiated !!')
_shutdown.set()
ws.close()
time.sleep(2)