From d8ab4d1d37e02b6a5743f8bbaa141aaf9065bf3e Mon Sep 17 00:00:00 2001
From: richardasaurus <richard@richard.do>
Date: Fri, 8 Feb 2013 12:21:32 +0000
Subject: [PATCH] Added get_link() and removed file['k'] from process_file()

Added get_link() for generating public file download link, changes by
jlejeune
---
 mega.py       | 29 +++++++++++++++++++++++++----
 tests/test.py | 13 +++++++------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/mega.py b/mega.py
index 34c15fd..c17e861 100644
--- a/mega.py
+++ b/mega.py
@@ -100,10 +100,31 @@ class Mega(object):
             file = file['f'][0]
             public_handle = self.api_request({'a': 'l', 'n': file['h']})
             file_key = file['k'][file['k'].index(':') + 1:]
-            decrypted_key = a32_to_base64(decrypt_key(base64_to_a32(file_key), self.master_key))
-            return '{0}://{1}/#!{2}!{3}'.format(self.schema, self.domain, public_handle, decrypted_key)
+            decrypted_key = a32_to_base64(decrypt_key(base64_to_a32(file_key),
+                                                      self.master_key))
+            return '{0}://{1}/#!{2}!{3}'.format(self.schema,
+                                                self.domain,
+                                                public_handle,
+                                                decrypted_key)
         else:
-            raise ValueError('This function requires upload() file object')
+            raise ValueError('Upload() response required as input, use get_link() for regular file input')
+
+    def get_link(self, file):
+        '''
+        Get a file public link from given file object
+        '''
+        file = file[1]
+        if 'h' in file and 'k' in file:
+            public_handle = self.api_request({'a': 'l', 'n': file['h']})
+            file_key = file['k'][file['k'].index(':') + 1:]
+            decrypted_key = a32_to_base64(decrypt_key(base64_to_a32(file_key),
+                                                      self.master_key))
+            return '{0}://{1}/#!{2}!{3}'.format(self.schema,
+                                                self.domain,
+                                                public_handle,
+                                                decrypted_key)
+        else:
+            raise errors.ValidationError('File id and key must be present')
 
     def download_url(self, url):
         path = self.parse_url(url).split('!')
@@ -316,7 +337,7 @@ class Mega(object):
             key = file['k'][file['k'].index(':') + 1:]
             key = decrypt_key(base64_to_a32(key), self.master_key)
             if file['t'] == 0:
-                k = file['k'] = (key[0] ^ key[4], key[1] ^ key[5], key[2] ^ key[6], key[3] ^ key[7])
+                k = (key[0] ^ key[4], key[1] ^ key[5], key[2] ^ key[6], key[3] ^ key[7])
                 file['iv'] = key[4:6] + (0, 0)
                 file['meta_mac'] = key[6:8]
             else:
diff --git a/tests/test.py b/tests/test.py
index 01e8ef7..d2308f5 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -2,8 +2,8 @@ from mega import Mega
 
 def test():
     #user details
-    email = 'your@email.com'
-    password = 'password'
+    email = 'richard@richard.do'
+    password = 'g3tmein1988'
 
     mega = Mega()
 
@@ -25,16 +25,17 @@ def test():
     print(m.upload('test.py'))
 
     ##get file's public link
-    #NOTE: currently this only works with upload() file obj, as below
-    file = m.upload('test.py')
-    print(m.get_upload_link(file))
+    #NOTE: if passing upload() function response use get_upload_link()
+    file = m.find('test.py')
+    #print(m.get_upload_link(file))
+    print(m.get_link(file))
 
     ##trash a file, by id or url
     #print(m.delete('f14U0JhD'))
     #print(m.delete_url('https://mega.co.nz/#!f14U0JhD!S_2k-EvB5U1N3s0vm3I5C0JN2toHSGkVf0UxQsiKZ8A'))
 
     ##search for a file in account
-    file = m.find('test.py')
+    file = m.find('somefile.doc')
     if file:
         #trash a file by it's id
         print(m.delete(file[1]['k']))