client.py

#!/usr/bin/env python3

import argparse
import base64
import json
import hashlib
import os
import stat

import requests

class Client:
    def __init__(self):
        self.http = requests.Session()
        self.http.hooks['response'].append(lambda r, *args, **kwargs : r.raise_for_status())
    
    def authenticate(self, token):
        self.http.headers['Authorization'] = "Bearer {}".format(token)
    
    def url(self, path=''):
        return "https://schedshare.herokuapp.com/{}".format(path)
    
    def request_token(self, user_id):
        return self.http.post(self.url("users/{}/token".format(user_id)))
    
    def get_binaries(self):
        return self.http.get(self.url('binaries')).json()
    
    def send_binary(self, filename, **kwargs):
        data = {'binary[{}]'.format(key): value for key, value in kwargs.items()}
        upload = ('binary[file]', ('upload.bin', open(filename, 'rb'), 'application/octet-stream'))
        self.http.post(self.url('binaries'), data=data, files=[upload])

class Config:
    def __init__(self, filename='.client.json'):
        self.filename = filename
        self.config = {}
        self.read_or_create()
    
    def read_or_create(self):
        try:
            with open(self.filename) as f:
                self.config = json.load(f)
        except FileNotFoundError:
            self.save()
    
    def update(self, entries={}, **kwargs):
        self.config.update(entries)
        self.config.update(kwargs)
        self.save()
    
    def __getitem__(self, key):
        return self.config[key]
    
    def __setitem__(self, key, value):
        self.update({key: value})
    
    def __contains__(self, key):
        return key in self.config

    def save(self):
        with open(self.filename, 'w') as f:
            json.dump(self.config, f)

class Downloader:
    def __init__(self, binaries, target):
        self.binaries = binaries
        self.target = target
    
    def checksum(self, filename):
        md5 = hashlib.md5()
        with open(filename, 'rb') as f:
            for chunk in iter(lambda: f.read(4096), b''):
                md5.update(chunk)
        return base64.b64encode(md5.digest()).decode('utf-8')
    
    def pull(self):
        for binary in self.binaries:
            filename = os.path.join(self.target, str(binary['user_id']) + '.exe')
            try:
                existing_checksum = self.checksum(filename)
            except FileNotFoundError:
                existing_checksum = None

            if existing_checksum == binary['checksum']:
                print("{} up to date".format(filename))
                continue
            
            download_file(binary['url'], filename)
            
            new_checksum = self.checksum(filename)
            if new_checksum == binary['checksum']:
                if existing_checksum is None:
                    os.chmod(filename, stat.S_IXUSR | os.stat(filename).st_mode)
                    print("{} downloaded".format(filename))
                else:
                    print("{} updated".format(filename))
            else:
                print("{} error".format(filename))
                os.remove(filename)

def download_file(url, destination):
    r = requests.get(url, stream=True)
    r.raise_for_status()
    with open(destination, 'wb') as f:
        for chunk in r:
            f.write(chunk)

client = Client()
config = Config()

if 'token' not in config:
    user_id = input('Enter your user ID: ')
    client.request_token(user_id)
    print('Access token has been sent to your email address.')
    token = input('Enter your access token: ')
    config['token'] = token

if 'target' not in config:
    target = input('Choose a destination folder for the binaries: ') or '.'
    real_target = os.path.realpath(target)
    os.makedirs(real_target, exist_ok=True)
    print('Binaries will be saved to {}'.format(real_target))
    config['target'] = real_target

client.authenticate(config['token'])

parser = argparse.ArgumentParser()
parser.add_argument('--upload', help='upload a binary file')
parser.add_argument('--self-update', help='update this script', action='store_true')
parser.add_argument('-v', help='version of an uploaded file')
args = parser.parse_args()

if args.self_update:
    script_url = 'https://gist.githubusercontent.com/hejmsdz/d0113e6cf89ae6d7e1997355c8477318/raw/client.py'
    download_file(script_url, __file__)
    print('updated successfully')
elif args.upload:
    client.send_binary(args.upload, version=args.v)
    print('uploaded successfully')
else:
    downloader = Downloader(client.get_binaries(), config['target'])
    downloader.pull()