Added File objects, some fixes, and download module

scripts/slack-backup

@@ -28,12 +28,15 @@ def main():
     parser.add_argument('-c', '--channels', default=[], nargs='+',
                         help='List of channels to perform actions on. '
                         'Default is all channels')
+    parser.add_argument('-t', '--team', default='', help='team name, which is'
+                        ' part of slack url, for example: if url is '
+                        '"https://team.slack.com" than "team" is a name of '
+                        'the team.')
     args = parser.parse_args()
 
-    slack = client.Client(args.token, args.database)
+    slack = client.Client(args)
-    slack.update_users()
+    slack.update()
-    slack.update_channels()
+    # slack.generate_history(Reporter(args.format))
-    slack.update_history(selected_channels=args.channels)
 
 
 if __name__ == "__main__":

slack_backup/client.py

@@ -1,21 +1,49 @@
 """
 Create backup for certain date for specified channel in slack
 """
-import logging
 from datetime import datetime
+import getpass
+import logging
 
 import slackclient
 
 from slack_backup import db
 from slack_backup import objects as o
+from slack_backup import download
 
 
 class Client(object):
-    def __init__(self, token, dbfilename=None):
+    """
-        self.slack = slackclient.SlackClient(token)
+    This class is intended to provide an interface for getting, storing and
-        self.engine = db.connect(dbfilename)
+    querying data fetched out using Slack API.
+    """
+    def __init__(self, args):
+        self.slack = slackclient.SlackClient(args.token)
+        self.engine = db.connect(args.dbfilename)
         self.session = db.Session()
+        self.selected_channels = args.channels
+        self.user = args.user
+        self.password = args.password
+        if not self.user and not self.password:
+            logging.warning('No media will be downloaded, due to not '
+                            'providing credentials for a slack account')
+        elif not self.user and self.password:
+            logging.warning('No media will be downloaded, due to not '
+                            'providing username for a slack account')
+        elif self.user and not self.password:
+            self.password = getpass.getpass(prompt='Provide password for '
+                                            'your slack account: ')
         self.q = self.session.query
+        self.dld = download.Download(args.user, args.password, args.team)
+
+    def update(self):
+        """
+        Perform an update, store data to db
+        """
+        self.dld.authorize()
+        self.update_users()
+        self.update_channels()
+        self.update_history()
 
     def update_channels(self):
         """Fetch and update channel list with current state in db"""
@@ -24,7 +52,7 @@ class Client(object):
         if not result:
             return
 
-        for channel_data in result['channels']:
+        for channel_data in result:
             channel = self.q(o.Channel).\
                 filter(o.Channel.slackid == channel_data['id']).one_or_none()
 
@@ -57,19 +85,19 @@ class Client(object):
 
         self.session.commit()
 
-    def update_history(self, selected_channels=None):
+    def update_history(self):
         """
         Get the latest or all messages out of optionally selected channels
         """
 
-        channels = self.q(o.Channel).all()
+        all_channels = self.q(o.Channel).all()
-        if selected_channels:
+        if self.selected_channels:
-            selected_channels = [c for c in channels
+            channels = [c for c in all_channels
-                                 if c.name in selected_channels]
+                        if c.name in self.selected_channels]
         else:
-            selected_channels = channels
+            channels = all_channels
 
-        for channel in selected_channels:
+        for channel in channels:
             latest = self.q(o.Message).\
                 filter(o.Message.channel == channel).\
                 order_by(o.Message.ts.desc()).first()
@@ -87,6 +115,10 @@ class Client(object):
         self.session.commit()
 
     def _create_message(self, data, channel):
+        """
+        Create message with corresponding possible metadata, like reactions,
+        files etc.
+        """
         message = o.Message(data)
         message.user = self.q(o.User).\
             filter(o.User.slackid == data['user']).one()
@@ -94,9 +126,17 @@ class Client(object):
 
         if 'reactions' in data:
             for reaction_data in data['reactions']:
-                o.Message.reactions.append(o.Reaction(reaction_data))
+                message.reactions.append(o.Reaction(reaction_data))
 
-        self.session.add(o.Message)
+        if data.get('subtype') == 'file_share':
+            message.file = o.File()
+            if data['file']['is_external']:
+                message.file.url = data['file']['url_private']
+            else:
+                priv_url = data['file']['url_private_download']
+                message.file.url = self.dld.get_local_url(priv_url)
+
+        self.session.add(message)
 
     def _get_create_obj(self, data, classobj, channel):
         """
@@ -162,7 +202,7 @@ class Client(object):
             logging.error(result['error'])
             return None
 
-        return result['channels']
+        return result['members']
 
     def _channels_history(self, channel, latest):
         """

slack_backup/download.py

@@ -0,0 +1,59 @@
+"""
+Module for download files, store them in local filesystem and convert the URLs
+to local ones, so that sophisticated writers can make a use of it
+"""
+import logging
+
+import requests
+
+
+class NotAuthorizedError(requests.HTTPError):
+    pass
+
+
+class Download(object):
+    """Download class for taking care of Slack internally uploaded files"""
+
+    def __init__(self, user, password, team):
+        self.session = requests.session()
+        self.team = team
+        self.user = user
+        self.password = password
+
+    def get_local_url(self, url):
+        """
+        Download file from provided url and save it locally. Return local URI.
+        """
+        # TODO: implementation
+        # res = session.post(url)
+        # new_path = self.prepare_uri(url)
+        # with open(new_path, "wb") as fobj:
+        #     fobj.write(p.content)
+        # return url
+
+        return url
+
+    def authorize(self):
+        """
+        Authenticate and gather session for Slack
+        """
+        res = self.session.get('https://%s.slack.com/' % self.team)
+
+        crumb = ''
+        for line in res.text.split('\n'):
+            if 'crumb' in line:
+                crumb = line.split('value=')[1].split('"')[1]
+                break
+        else:
+            logging.error('Cannot access Slack login page')
+            raise NotAuthorizedError('Cannot access Slack login page')
+
+        res = self.session.post("https://%s.slack.com/" % self.team,
+                                {'crumb': crumb,
+                                 'email': self.user,
+                                 'password': self.password,
+                                 'signin': 1})
+        cookies = requests.utils.dict_from_cookiejar(self.session.cookies)
+        if not ('a' in cookies and 'b' in cookies and
+                ('a-' + cookies['a']) in cookies):
+            raise NotAuthorizedError('Failed to login into Slack app')

slack_backup/objects.py

@@ -219,6 +219,7 @@ class Message(Base):
     channel = relationship("Channel", back_populates="messages")
 
     reactions = relationship("Reaction", back_populates="message")
+    files = relationship("File", back_populates="message")
 
     def __init__(self, data_dict=None):
         self.update(data_dict)
@@ -232,3 +233,15 @@ class Message(Base):
         self.ts = float(data_dict.get('ts', 0))
         self.text = data_dict.get('text', '')
         self.type = data_dict.get('subtype', '')
+
+
+class File(Base):
+    __tablename__ = "files"
+
+    id = Column(Integer, primary_key=True)
+    url = Column(Text)
+    thumbnail = Column(Text)
+    relative_path = Column(Text)
+
+    message_id = Column(Integer, ForeignKey('messages.id'))
+    message = relationship('Message', back_populates='files')

tests/test_client.py

@@ -1,8 +1,8 @@
-import unittest
+from unittest import TestCase
 from unittest.mock import MagicMock
 
 from slack_backup import client
-from slack_backup import objects
+from slack_backup import objects as o
 
 CHANNELS = {"ok": True,
             "channels": [{"id": "C00000000",
@@ -177,12 +177,6 @@ USERS = {'cache_ts': 1479577519,
                       'tz_label': 'Pacific Standard Time',
                       'tz_offset': -28800}]}
 
-MSG2 = {"type": "message",
-        "user": "UCCCCCCCC",
-        "text": "Pellentesque molestie nunc id enim. Etiam mollis tempus "
-                "neque. Duis. per conubia nostra, per",
-        "ts": "1479505026.000002"}
-
 MSGS = {'messages': [{"type": "message",
                       "user": "UAAAAAAAA",
                       "text": "Class aptent taciti sociosqu ad litora torquent"
@@ -333,48 +327,116 @@ MSGS = {'messages': [{"type": "message",
                       "upload": True}],
         "ok": True,
         "latest": "1479501075.000020",
-        "has_more": False}
+        "has_more": True}
+
+MSG2 = {'messages': [{"type": "message",
+                      "user": "UCCCCCCCC",
+                      "text": "Pellentesque molestie nunc id enim. Etiam "
+                              "mollis tempus neque. Duis. per conubia "
+                              "nostra, per",
+                      "ts": "1479505026.000002"}],
+        "ok": True,
+        "latest": "1479505026.000003",
+        "has_more": True}
+
+MSG3 = {"ok": True,
+        "oldest": "1479505026.000003",
+        "messages": [],
+        "has_more": False,
+        "is_limited": False}
 
 
-class TestApiCalls(unittest.TestCase):
+class FakeArgs(object):
+    token = 'token_string'
+    user = 'fake_user'
+    password = 'fake_password'
+    team = 'fake_team'
+    dbfilename = None
+    channels = None
 
-    def setup(self):
+
-        print("asd")
+class TestApiCalls(TestCase):
 
     def test_channels_list(self):
-        self.assertTrue(1)
+        cl = client.Client(FakeArgs())
+        cl.slack.api_call = MagicMock(return_value=CHANNELS)
+        channels = cl._channels_list()
+        self.assertListEqual(CHANNELS['channels'], channels)
 
     def test_users_list(self):
-        self.assertTrue(1)
+        cl = client.Client(FakeArgs())
+        cl.slack.api_call = MagicMock(return_value=USERS)
+        users = cl._users_list()
+        self.assertListEqual(USERS['members'], users)
 
     def test_channels_history(self):
-        self.assertTrue(1)
+        cl = client.Client(FakeArgs())
 
-
-class TestClient(unittest.TestCase):
-
-    def test_update_users(self):
-        cl = client.Client("token string")
         cl.slack.api_call = MagicMock(return_value=USERS)
         cl.update_users()
-        users = cl.session.query(objects.User).all()
+
+        cl.slack.api_call = MagicMock(return_value=CHANNELS)
+        cl.update_channels()
+
+        cl.slack.api_call = MagicMock()
+        cl.slack.api_call.side_effect = [MSGS, MSG2, MSG3]
+
+        channel = cl.q(o.Channel).filter(o.Channel.slackid ==
+                                         "C00000001").one()
+
+        msg, ts = cl._channels_history(channel, 0)
+        self.assertEqual(len(msg), 5)
+        self.assertEqual(ts, '1479501074.000032')
+
+        msg, ts = cl._channels_history(channel, ts)
+        self.assertEqual(len(msg), 1)
+        self.assertEqual(ts, '1479505026.000002')
+
+        msg, ts = cl._channels_history(channel, ts)
+        self.assertEqual(len(msg), 0)
+        self.assertIsNone(ts)
+
+
+class TestClient(TestCase):
+
+    def test_update_users(self):
+        cl = client.Client(FakeArgs())
+        cl.slack.api_call = MagicMock(return_value=USERS)
+        cl.update_users()
+        users = cl.session.query(o.User).all()
         self.assertEqual(len(users), 4)
         self.assertEqual(users[0].id, 1)
 
         cl.update_users()
-        users = cl.session.query(objects.User).all()
+        users = cl.session.query(o.User).all()
         self.assertEqual(len(users), 4)
         self.assertEqual(users[0].id, 1)
         self.assertEqual(users[0].slackid, 'UAAAAAAAA')
 
 
-class TestMessage(unittest.TestCase):
+class TestMessage(TestCase):
 
     def setUp(self):
-        self.cl = client.Client('token string')
+        args = FakeArgs()
+        args.channels = ['general']
+
+        self.cl = client.Client(args)
+        self.cl.dld.authorize = MagicMock()
+        self.cl.slack.api_call = MagicMock(return_value=USERS)
+        self.cl.update_users()
+
+        self.cl.slack.api_call = MagicMock(return_value=CHANNELS)
+        self.cl.update_channels()
+
         self.cl.slack.api_call = MagicMock()
 
-    def test_create_message(self):
+    def test_update_history(self):
 
-        cl = client.Client("token string")
+        self.cl.slack.api_call.side_effect = [MSGS, MSG3]
-        cl.slack.api_call = MagicMock(return_value=MSGS)
+        self.cl.update_history()
+        self.assertEqual(len(self.cl.q(o.Message).all()), 5)
+
+        self.cl.slack.api_call.side_effect = [MSG2, MSG3]
+        self.cl.update_history()
+
+        self.assertEqual(len(self.cl.q(o.Message).all()), 6)