megan.py

   1 #! /usr/bin/env python3
   2
   3 from urllib.parse import urlsplit
   4 from base64 import urlsafe_b64decode
   5 from net import request_cached
   6 import json
   7 import subprocess
   8 import urllib.request
   9 from shutil import copyfileobj
  10 import selectors
  11 from shorthand import bitmask
  12 from contextlib import ExitStack
  13 import os
  14 import hmac
  15 from io import SEEK_CUR
  16 from datacopy import TerminalLog, Progress
  17 from io import TextIOWrapper
  18 from enum import Enum
  19 import gzip
  20 from net import header_list
  21
  22 def main(url, start='0', *, v=False, n=False):
  23     megan = Megan(url)
  24     API = 'https://g.api.mega.co.nz/'
  25     if megan.folder is None:
  26         [url, headers, body] = megan.get_file_request()
  27     else:
  28         [url, headers, body] = megan.get_folder_request()
  29         with ExitStack() as cleanup:
  30             [header, resp] = request_cached(API + url,
  31                 headers=headers, types=('application/json',),
  32                 data=body, cleanup=cleanup)
  33             megan.handle_folder_response(header, resp.read())
  34
  35         [megan.key, req] = megan.get_folder_file_request(megan.file)
  36         [url, headers, body] = req
  37     with ExitStack() as cleanup:
  38         [header, resp] = request_cached(API + url,
  39             headers=headers, types=('application/json',),
  40             data=body, cleanup=cleanup)
  41         [K, b] = megan.handle_file_response(header, resp.read())
  42     print('Encryption key:', K)
  43
  44     print('Filename:', megan.name)
  45     print('Size: {}'.format(format_size(megan.resp['s'])))
  46
  47     start = int(start, 0)
  48     if start % 16:
  49         start -= start % 16
  50         print('Starting on block boundary at +0x{:X}'.format(start))
  51     if v:
  52         print('Download URL:', megan.resp['g'])
  53
  54     print('IV:', megan.key[16:24].hex().upper())
  55     print('MAC:', megan.key[24:32].hex().upper())
  56     if n:
  57         return
  58
  59     req = urllib.request.Request(megan.resp['g'])
  60     with ExitStack() as cleanup:
  61         if start:
  62             name = '{}+0x{:X}'.format(megan.name, start)
  63             print('Writing to', name)
  64             req.add_header('Range', 'bytes={}-'.format(start))
  65             download = cleanup.enter_context(urllib.request.urlopen(req))
  66             file = cleanup.enter_context(open(name, 'xb'))
  67             iv = (b & ~0 << 64) + (start // 16 & bitmask(128))
  68             decryptor = make_encryptor('-aes-128-ctr', '-d', K, iv,
  69                 stdout=file)
  70             cleanup.enter_context(decryptor)
  71             file.close()
  72             copyfileobj(download, decryptor.stdin)
  73             assert not decryptor.wait()
  74         else:
  75             chunk_iv = b & ~0 << 64 | b >> 64
  76             authenticator = Authenticator(megan.resp['s'], K, chunk_iv)
  77             cleanup.callback(authenticator.close)
  78             try:
  79                 file = cleanup.enter_context(open(megan.name, 'xb'))
  80             except FileExistsError:
  81                 file = cleanup.enter_context(open(megan.name, 'r+b'))
  82                 while True:
  83                     data = file.read(0x10000)
  84                     if len(data) < 0x10000:
  85                         break
  86                     authenticator.update(data)
  87                 if authenticator.pos + len(data) == megan.resp['s']:
  88                     authenticator.update(data)
  89                 else:
  90                     limit = len(data) - len(data) % 16
  91                     with memoryview(data) as view, view[:limit] as slice:
  92                         authenticator.update(slice)
  93                     file.seek(-len(data) % 16, SEEK_CUR)
  94                     pos = authenticator.pos
  95                     print('Continuing file from {}'.format(format_size(pos)))
  96                     req.add_header('Range', 'bytes={}-'.format(pos))
  97             if authenticator.pos < megan.resp['s']:
  98                 download = cleanup.enter_context(urllib.request.urlopen(req))
  99                 iv = (b & ~0 << 64) + (authenticator.pos // 16 & bitmask(128))
 100                 decryptor = make_encryptor('-aes-128-ctr', '-d', K, iv,
 101                     stdout=subprocess.PIPE)
 102                 cleanup.enter_context(decryptor)
 103                 selector = cleanup.enter_context(selectors.DefaultSelector())
 104
 105                 download_buf = b''
 106                 download_pos = 0x10000
 107
 108                 log = TerminalLog()
 109                 cleanup.callback(Progress.close, log)
 110                 progress = Progress(log, megan.resp['s'])
 111                 total_dl = 0
 112
 113                 def copy_download():
 114                     nonlocal download_buf, download_pos
 115                     try:
 116                         if download_pos >= len(download_buf):
 117                             if download_pos < 0x10000:
 118                                 raise EOFError()
 119                             download_buf = download.read(0x10000)
 120                             if not download_buf:
 121                                 raise EOFError()
 122                             download_pos = 0
 123                     except EOFError:
 124                         selector.unregister(decryptor.stdin)
 125                         decryptor.stdin.close()
 126                     else:
 127                         with memoryview(download_buf) as view, \
 128                                 view[download_pos:] as rest:
 129                             n = decryptor.stdin.write(rest)
 130                         download_pos += n
 131
 132                         nonlocal total_dl
 133                         total_dl += n
 134                         progress.update(total_dl)
 135                 os.set_blocking(decryptor.stdin.fileno(), False)
 136                 selector.register(decryptor.stdin, selectors.EVENT_WRITE,
 137                     copy_download)
 138
 139                 def read_data():
 140                     data = decryptor.stdout.read(0x10000)
 141                     if not data:
 142                         raise EOFError()
 143                     file.write(data)
 144                     authenticator.update(data)
 145                 selector.register(decryptor.stdout, selectors.EVENT_READ,
 146                     read_data)
 147
 148                 while authenticator.pos < megan.resp['s']:
 149                     for [k, events] in selector.select():
 150                         k.data()
 151             assert hmac.compare_digest(authenticator.digest(), megan.key[24:])
 152             print('MAC validated')
 153
 154 class Megan:
 155     def __init__(self, url):
 156         split = urlsplit(url)
 157         assert split.netloc == 'mega.nz'
 158         assert split.scheme == 'https'
 159         assert not split.query
 160
 161         if split.path == '/':
 162             self.folder = None
 163             assert split.fragment.startswith('!')
 164             [self.id, self.key] = split.fragment[1:].split('!', 2)
 165             self.key = base64_decode(self.key, 256)
 166             return
 167
 168         assert split.path.startswith('/folder/')
 169         self.folder = split.path[len('/folder/'):]
 170         self.key = split.fragment.split('/file/', 2)
 171         if len(self.key) == 1:
 172             [self.key] = self.key
 173             self.file = None
 174         else:
 175             [self.key, self.file] = self.key
 176         self.key = base64_decode(self.key, 128).hex()
 177         return
 178
 179     def get_folder_request(self):
 180         return make_request('f', {
 181             'r': 1,  # Recursive
 182             'ca': 1,  # Cache
 183         }, f'?n={self.folder}')
 184
 185     def handle_folder_response(self, *folder):
 186         folder = decode_response(*folder)
 187         self.nodes = dict()
 188         for f in folder['f']:
 189             node = {
 190                 'key': base64_decode(f['k'].split(':', 1)[1]),
 191                 'a': f['a'],
 192                 'type': NodeType(f['t']),
 193                 'parent': f['p'],
 194                 'ts': f['ts'],
 195             }
 196             existing = self.nodes.setdefault(f['h'], node)
 197             assert existing is node
 198             if node['type'] is NodeType.FILE:
 199                 node['size'] = f['s']
 200         self.root = folder['f'][0]['h']
 201         node = self.nodes[self.root]
 202         [k, b] = split_key(self.decrypt_key(node))
 203         self.name = decrypt_name(node['a'], k)
 204
 205     def get_folder_file_request(self, file):
 206         key = self.decrypt_key(self.nodes[file])
 207
 208         args = {
 209             'n': file,
 210             'g': 1,  # Include download URL
 211         }
 212         return (key, make_request('g', args, f'?n={self.folder}'))
 213
 214     def get_file_request(self):
 215         args = {
 216             'p': self.id,
 217             'g': 1,  # Include download URL
 218         }
 219         return make_request('g', args)
 220
 221     def handle_file_response(self, *resp):
 222         self.resp = decode_response(*resp)
 223         self.root = None
 224         self.nodes = {None: {
 225             'type': NodeType.FILE,
 226             'a': self.resp['at'],
 227             'size': self.resp['s'],
 228         }}
 229
 230         [K, b] = split_key(self.key)
 231         self.name = decrypt_name(self.resp['at'], K)
 232         return (K, b)
 233
 234     def decrypt_key(self, node):
 235         with make_encryptor('-aes-128-ecb', '-d', self.key,
 236                 stdout=subprocess.PIPE) as dec:
 237             [k, err] = dec.communicate(node['key'])
 238         assert dec.returncode == 0
 239         return k
 240
 241 def split_key(key):
 242     a = int.from_bytes(key[:16], 'big')
 243     b = int.from_bytes(key[16:], 'big')
 244     return (format(a ^ b, '032X'), b)
 245
 246 def decrypt_name(a, k):
 247     dec = make_encryptor('-aes-128-cbc', '-d', k, 0, stdout=subprocess.PIPE)
 248     with dec:
 249         [a, err] = dec.communicate(base64_decode(a))
 250     assert dec.returncode == 0
 251     assert a.startswith(b'MEGA')
 252     a = a[4:].rstrip(b'\x00')
 253     a = json.loads(a.decode('utf-8'))
 254     return a['n']
 255
 256 class NodeType(Enum):
 257     FILE = 0
 258     FOLDER = 1
 259
 260 def make_request(type, args, url=''):
 261     args['a'] = type
 262     return ('cs' + url, (
 263         ('User-Agent', 'megan'),
 264         ('Content-Type', 'application/json'),
 265         ('Accept-Encoding', 'gzip, x-gzip'),
 266     ), json.dumps((args,)).encode('ascii'))
 267
 268 def decode_response(header, resp):
 269     coded = False
 270     for code in header_list(header, 'Content-Encoding'):
 271         if code.lower() in {'gzip', 'x-gzip'}:
 272             if coded:
 273                 raise TypeError('Recursive gzip encoding')
 274             resp = gzip.decompress(resp)
 275             coded = True
 276         else:
 277             raise TypeError(f'Unhandled encoding: {code!r}')
 278
 279     resp = json.loads(resp)
 280     if not isinstance(resp, list):
 281         error(resp)
 282     [resp] = resp
 283     if not isinstance(resp, dict):
 284         error(resp)
 285     return resp
 286
 287 ERRORS = {
 288     -1: 'INTERNAL',
 289     -2: 'ARGS',
 290     -3: 'AGAIN: Temporary congestion',
 291     -4: 'RATELIMIT',
 292     -5: 'FAILED',
 293     -6: 'TOOMANY: Connection limit exceeded or terms breached',
 294     -7: 'RANGE',
 295     -8: 'EXPIRED',
 296     -9: 'NOENT',
 297     -10: 'CIRCULAR',
 298     -11: 'ACCESS',
 299     -12: 'EXIST',
 300     -13: 'INCOMPLETE',
 301     -14: 'KEY: Decryption failed',
 302     -15: 'SID: Invalid session',
 303     -16: 'BLOCKED',
 304     -17: 'OVERQUOTA',
 305     -18: 'TEMPUNAVAIL',
 306     -19: 'TOOMANYCONNECTIONS',
 307     -20: 'WRITE',
 308     -21: 'READ',
 309     -22: 'APPKEY',
 310 }
 311 def error(code):
 312     raise Exception(f'{ERRORS[code]} ({code})')
 313
 314 def base64_decode(s, bits=None):
 315     if bits is not None:
 316         assert len(s) == ceildiv(bits, 6)
 317     return urlsafe_b64decode(s + '=' * (-len(s) % 4))
 318
 319 class Authenticator(ExitStack):
 320     def __init__(self, size, K, chunk_iv):
 321         ExitStack.__init__(self)
 322         self.size = size
 323         self.K = K
 324         self.whole_mac = self.enter_context(CbcMac(self.K))
 325         self.pos = 0
 326         self.chunk_size = 0
 327         self.chunk_mac = None
 328         self.chunk_cleanup = self.enter_context(ExitStack())
 329         assert self.size > 0
 330         self.padding = -self.size % 16
 331         self.end = self.size + self.padding
 332         self.chunk_iv = chunk_iv
 333
 334     def update(self, data):
 335         with memoryview(data) as view:
 336             pos = 0
 337             while True:
 338                 if not self.chunk_mac:
 339                     # New MAC chunk
 340                     self.chunk_mac = CbcMac(self.K, self.chunk_iv)
 341                     self.chunk_cleanup.callback(self.chunk_mac.close)
 342                     if self.chunk_size < 1 << 20:
 343                         self.chunk_size += 0x20000
 344                     self.chunk_pos = 0
 345                 chunk_end = pos + min(self.chunk_size - self.chunk_pos,
 346                     self.end - self.pos)
 347                 with view[pos:max(chunk_end - 16, 0)] as early:
 348                     self.chunk_mac.write(early)
 349                     pos += len(early)
 350                     self.chunk_pos += len(early)
 351                     self.pos += len(early)
 352                 if pos < chunk_end - 16:
 353                     break
 354                 block = view[pos:chunk_end]
 355                 self.chunk_mac.write_last_block(block)
 356                 pos += len(block)
 357                 self.chunk_pos += len(block)
 358                 self.pos += len(block)
 359                 if self.pos == self.size:
 360                     self.chunk_mac.write_last_block(bytes(self.padding))
 361                     self.whole_mac.write_last_block(self.chunk_mac.get_mac())
 362                     return
 363                 if pos < chunk_end:
 364                     break
 365                 self.whole_mac.write(self.chunk_mac.get_mac())
 366                 self.chunk_cleanup.close()
 367                 self.chunk_mac = None
 368
 369     def digest(self):
 370         mac = int.from_bytes(self.whole_mac.get_mac(), 'big')
 371         mac ^= mac >> 32
 372         mac = mac.to_bytes(16, 'big')
 373         return mac[4:8] + mac[12:16]
 374
 375 class CbcMac(ExitStack):
 376     def __init__(self, K, iv=0):
 377         ExitStack.__init__(self)
 378         self.proc = make_encryptor('-aes-128-cbc', '-e', K, iv,
 379             stdout=subprocess.PIPE)
 380         self.enter_context(self.proc)
 381         self.selector = self.enter_context(selectors.DefaultSelector())
 382         os.set_blocking(self.proc.stdin.fileno(), False)
 383         self.selector.register(self.proc.stdin, selectors.EVENT_WRITE,
 384             self.feed_data)
 385         self.selector.register(self.proc.stdout, selectors.EVENT_READ,
 386             self.skip_enc)
 387         self.to_skip = 0
 388
 389     def write(self, data):
 390         self.to_skip += len(data)
 391         self.write_last_block(data)
 392
 393     def feed_data(self):
 394         with self.data[self.data_pos:] as data:
 395             self.data_pos += self.proc.stdin.write(data)
 396
 397     def skip_enc(self):
 398         chunk = min(self.to_skip, 0x10000)
 399         self.to_skip -= len(self.proc.stdout.read(chunk))
 400
 401     def write_last_block(self, data):
 402         with memoryview(data) as self.data:
 403             self.data_pos = 0
 404             while self.data_pos < len(data):
 405                 for [key, events] in self.selector.select():
 406                     key.data()
 407
 408     def get_mac(self):
 409         self.proc.stdin.close()
 410         while self.to_skip:
 411             self.skip_enc()
 412         return self.proc.stdout.readall()
 413
 414 def make_encryptor(cipher, dir, K, iv=None, *, stdout):
 415     args = ['openssl', 'enc', cipher, '-nopad', dir, '-K', K]
 416     if iv is not None:
 417         args.extend(('-iv', format(iv, '032X')))
 418     return subprocess.Popen(args,
 419         stdin=subprocess.PIPE, stdout=stdout, bufsize=0)
 420
 421 def format_size(s):
 422     si = format(s)
 423     [prefix, scaled_exp] = divmod(len(si) - 1, 3)
 424     if prefix:
 425         whole = si[:1 + scaled_exp]
 426         fraction = si[1 + scaled_exp:]
 427         prefix = 'kMGTPE'[prefix - 1]
 428         si = '{}.{} {}B'.format(whole, fraction, prefix)
 429         prefix = (s.bit_length() - 1) // 10
 430         bin = s / 1024**prefix
 431         digits = format(bin, '#.3g')
 432         if 'e' in digits:
 433             prefix += 1
 434             bin /= 1024
 435             digits = format(bin, '#.3g')
 436         prefix = 'KMGTPE'[prefix - 1]
 437         size = '{} ({} {}iB)'.format(si, digits.rstrip('.'), prefix)
 438     else:
 439         size = si + ' B'
 440     return size
 441
 442 def ceildiv(a, b):
 443     return -(-a // b)
 444
 445 if __name__ == "__main__":
 446     import clifunc
 447     clifunc.run()