fastkeybase

#!/usr/local/bin/python3
from sh import keybase
from argparse import ArgumentParser
from requests import get, post
from re import compile as rcompile
from os import getenv
from os.path import exists
from tempfile import NamedTemporaryFile
from subprocess import call
from json import dumps

def parse_args():
  p = ArgumentParser(description='A more streamlined keybase wrapper.')
  p.add_argument('-w', '--who', help='Who the message is being encrypted for.', dest='who')
  p.add_argument('-m', '--message', help='The message to encrypt, if encrypting. If no message is provided, $EDITOR will be opened to retrieve a message.', nargs='?', dest='message', const=False)
  p.add_argument('-f', '--file', help='The name of the file in the Gist to decrypt, if decrypting a Gist. Defaults to the first file.', dest='file')
  p.add_argument('-t', '--type', help='If decrypting, the type of data specified by "data." If not specified, a guess will be attempted. If encrypting, what to do with the encrypted data. Either raw option will only print out the raw encrypted message. If not specified, the encrypted data will be printed.', dest='type', choices=['gist', 'hastebin', 'raw-web', 'raw-file'])
  p.add_argument('-d', '--data', help='The data to decrypt.', dest='data')
  p.add_argument('-l', '--loud', help='Output more status lines.', dest='loud', default=False, action='store_true')
  return p.parse_args()

class FastKeybaseEncrypt:

  GIST_API_URL = 'https://api.github.com/gists'
  HASTEBIN_API_URL = 'http://hastebin.com/documents'

  def __init__(self, args):
    self.message = args['message']
    if self.message is None:
      print('No data provided.')
      exit(1)
      return
    self.who = args['who']
    if self.who is None:
      print('No username provided.')
      exit(1)
      return
    self.type = args['type'] if args['type'] is not None else 'raw-file'
    self.args = args

  def get_message(self):
    editor = "editor"
    with NamedTemporaryFile() as f:
      call(editor.split(' ') + [f.name])
      with open(f.name) as ff:
        data = ff.read()
    return data

  def check_for_encrypted_data(self):
    if self.encrypted is None:
      print('No encrypted data to Gist.')
      exit(1)
      return False
    return True

  def gist(self):
    if not self.check_for_encrypted_data():
      return
    j = post(self.GIST_API_URL, data=dumps({
      'files': {
        'file.txt': {
          'content': self.encrypted
        }
      }
    })).json()
    print(j['html_url'])

  def hastebin(self):
    if not self.check_for_encrypted_data():
      return
    j = post(self.HASTEBIN_API_URL, data=self.encrypted).json()
    print('http://hastebin.com/{}'.format(j['key']))

  def encrypt(self):
    if not self.message:
      self.log('Opening editor to get message...')
      self.message = self.get_message()
    self.encrypt_message()

  def encrypt_message(self):
    self.log('Calling keybase binary. Please wait...')
    self.encrypted = keybase('encrypt', self.who, m=self.message, batch=True).stdout.decode('utf-8')
    if self.type == 'raw-file' or self.type == 'raw-web':
      print(self.encrypted)
    elif self.type == 'gist':
      self.gist()
    elif self.type == 'hastebin':
      self.hastebin()

  def log(self, msg):
    if self.args['loud']:
      print(msg)

class FastKeybaseDecrypt:

  GIST_API_URL = 'https://api.github.com/gists/{}'
  GIST_PATTERN = rcompile(r'https?:\/\/gist\.github(usercontent)?\.com\/(\w+)(\/(\w+))?(\/raw\/\w+\/.+)?')

  HASTEBIN_API_URL = 'http://hastebin.com/raw/{}'
  HASTEBIN_PATTERN = rcompile(r'https?:\/\/(?!www\.)?hastebin\.com\/(raw\/)?(\w+)(?!.+)?')

  def __init__(self, args):
    self.data = args['data']
    if self.data is None:
      print('No data provided.')
      exit(1)
      return
    self.args = args

  def log(self, msg):
    if self.args['loud']:
      print(msg)

  def guess_type(self):
    if self.GIST_PATTERN.match(self.data):
      return 'gist'
    elif self.HASTEBIN_PATTERN.match(self.data):
      return 'hastebin'
    elif exists(self.data):
      return 'raw-file'
    return 'raw-web'

  def decrypt(self):
    decrypt_type = self.args['type'].lower() if self.args['type'] is not None else self.guess_type()
    if decrypt_type == 'gist':
      self.decrypt_gist()
    elif decrypt_type == 'hastebin':
      self.decrypt_hastebin()
    elif decrypt_type == 'raw-web':
      self.decrypt_raw_web()
    elif decrypt_type == 'raw-file':
      self.decrypt_raw_file()
    else:
      print('Could not guess type.')
      exit(1)

  def decrypt_gist(self):
    self.log('Decrypting Gist...')
    m = self.GIST_PATTERN.match(self.data)
    if not m:
      print('Invalid Gist URL.')
      exit(1)
      return
    if None not in m.group(1, 5):
      self.decrypt_raw_web()
      return
    gist_id = m.group(4) if None not in m.group(2, 3, 4) else m.group(2)
    if gist_id is None:
      print('Could not get Gist ID.')
      exit(1)
      return
    j = get(self.GIST_API_URL.format(gist_id)).json()
    f = args['file']
    fs = j['files']
    if len(fs) < 1:
      print('No files.')
      exit(1)
      return
    if f is not None and f not in fs:
      print('No such file.')
      exit(1)
      return
    if f is None:
      f = list(fs)[0]
    if f not in fs:
      print('Invalid file.')
      exit(1)
      return
    self.data = j['files'][f]['raw_url']
    self.decrypt_raw_web()

  def decrypt_hastebin(self):
    self.log('Decrypting hastebin...')
    m = self.HASTEBIN_PATTERN.match(self.data)
    if not m:
      print('Invalid hastebin URL.')
      exit(1)
      return
    hastebin_id = m.group(2)
    self.data = self.HASTEBIN_API_URL.format(hastebin_id)
    self.decrypt_raw_web()

  def decrypt_raw_web(self):
    self.log('Decrypting raw web...')
    self.data = get(self.data).text
    self.decrypt_raw()

  def decrypt_raw(self):
    self.log('Calling keybase binary. Prepare to enter password.')
    print(keybase('decrypt', m=self.data, batch=True))
    self.log('\nFinished!')

  def decrypt_raw_file(self):
    self.log('Decrypting raw file...')
    with open(self.data) as f:
      self.data = f.read()
    decrypt_raw()

def main(args):
  if args['message'] is not None:
    FastKeybaseEncrypt(args).encrypt()
  elif args['data'] is not None:
    FastKeybaseDecrypt(args).decrypt()
  else:
    print('Nothing to do.')
    exit(1)

if __name__ == '__main__':
  args = vars(parse_args())
  try:
    main(args)
  except Exception as e:
    print('An error occurred: {}'.format(e))
    exit(1)