romeotools/api.py

import json
import datetime
import sqlite3

from request import Request, urlopen

try:
    from settings import USERNAME, PASSWORD
except:
    USERNAME, PASSWORD = None, None

API_URL = "https://www.planetromeo.com/api/v4/"
IMAGE_URL = "https://www.planetromeo.com/img/usr/original/0x0/%s.jpg"

class Database:
    def __init__(self, db="database.db"):
        self.conn = sqlite3.connect(db)
        self.cur = self.conn.cursor()

    def create_tables(self):
        query = """
        CREATE TABLE IF NOT EXISTS `message` (`id` VARCHAR(64), `sender` VARCHAR(64), `recipient` VARCHAR(64), `date` TIMESTAMP, `text` TEXT);
        CREATE TABLE IF NOT EXISTS `attachment` (`id` VARCHAR(64) PRIMARY KEY, `owner` VARCHAR(64), `token` VARCHAR(64), `message` VARCHAR(64), `content` BLOB, FOREIGN KEY(`message`) REFERENCES message(`id`));
        CREATE TABLE IF NOT EXISTS `profile` (`id` VARCHAR(64), `name` VARCHAR(256));
        CREATE TABLE IF NOT EXISTS `location` (`profile` VARCHAR(64), `timestamp` TIMESTAMP, `name` VARCHAR(128), `country` VARCHAR(64), FOREIGN KEY(`profile`) REFERENCES profile(`id`));
        """

        for line in query.strip().split("\n"):
            self.cur.execute(line)

    def execute(self, *args, **kwargs):
        self.cur.execute(*args, **kwargs)
        self.conn.commit()

    def fetchone(self, *args, **kwargs):
        return self.cur.fetchone(*args, **kwargs)    

class Attachment:
    def __init__(self, id, owner, token, message, content=None, retrieve=True):
        self.id = id
        self.owner = owner
        self.token = token
        self.content = content if content else self.get_content() if retrieve else None
        self.message = message

    def get_content(self, exceptions=False):
        try:
            return urlopen(Request(IMAGE_URL % self.token)).read()
        except:
            if not exceptions:
                return False
            raise

    def to_database(self, exceptions=False):
        try:
            Database().execute("INSERT INTO attachment VALUES (?, ?, ?, ?, ?);", (self.id, self.owner, self.token, self.message, self.content))
        except:
            if not exceptions:
                return
            raise

    @classmethod
    def from_dict(cls, attachment, message_id):
        if attachment["type"] == "IMAGE":
            return cls(attachment["params"]["id"], attachment["params"]["owner_id"], attachment["params"]["url_token"], message_id)

class Location:
    def __init__(self, profile, name, country, timestamp=datetime.datetime.now()):
        self.profile = profile
        self.name = name
        self.country = country
        self.timestamp = timestamp

    def to_database(self, exceptions=False):
        try:
            Database().execute("INSERT INTO location VALUES (?, ?, ?, ?);", (self.profile, self.timestamp, self.name, self.country))
        except:
            if not exceptions:
                return
            raise

class Profile:
    def __init__(self, id, name):
        self.id = id
        self.name = name

    def to_database(self, exceptions=False):
        try:
            Database().execute("INSERT INTO profile VALUES (?, ?);", (self.id, self.name))
        except:
            if not exceptions:
                return
            raise

class Message:
    def __init__(self, sender, recipient, id, date, text, attachments=[]):
        self.sender = sender
        self.recipient = recipient
        self.id = id
        self.date = date
        self.text = text
        self.attachments = attachments

    def to_database(self, exceptions=False):
        try:
            Database().execute("INSERT INTO message VALUES (?, ?, ?, ?, ?);", (self.id, self.sender, self.recipient, self.date, self.text))
            for attachment in self.attachments:
                attachment.to_database()
        except:
            if not exceptions:
                return
            raise

    @classmethod
    def from_dict(cls, message):
        if not "attachments" in message.keys():
            message["attachments"] = []
        return cls(message["from"], message["to"], message["id"], message["date"], message["text"], [Attachment.from_dict(attachment, message["id"]) for attachment in message["attachments"] if attachment["type"] == "IMAGE"])


class RomeoAPI:
    def __init__(self, login=True, username=USERNAME, password=PASSWORD):
        if login:
            self.login(username, password)

    def login(self, username, password):
        payload = {"username": USERNAME, "password": PASSWORD, "keep_login": True}
        response = urlopen(Request(API_URL + "session?lang=en"), json.dumps(payload).encode("utf-8"))
        data = json.load(response)
        self.session = data

    def build_headers(self, additional={}):
        headers = {
            "X-Api-Key": "vuEp8o93b34CxUCljSMFEdhI70qDWtuk",
            "X-Session-Id": self.session["session_id"]
        }

        for key, value in additional:
            headers[key] = value

        return headers

    def messages(self, count=16384):
        headers = self.build_headers()
        response = urlopen(Request(API_URL + "messages?lang=en&length=%i" % count, headers=headers))
        data = json.load(response)
        
        for message in data["items"]:
            yield Message.from_dict(message)