#!/usr/bin/python3

import argparse
import base64
import os
import sys

BANNER = "SANTAS+SECRET+CHRISTMAS+CODE//USE+xmasdecode+TO+UNPACK\n"
SUFFIX = ".xmascode"

# reverses a string, = characters stay where they were
def reverse_equals(inp):
  out = ""
  for char in inp:
    if char == '=':
      out = out + char
    else:
      out = char + out

  return out

def strip_header(inp):
  return inp[len(BANNER):]

def prepend_header(inp):
  return BANNER + inp

# returns byte array xmas64-decoded from the inp
def xmasdecode(inp):
  # reverse the string
  inp = reverse_equals(inp)
  # feed it through base64
  out = base64.b64decode(inp.encode())
  return out

# returns xmas64-encoded string
def xmasencode(inp):
  # feed it through base64
  inp = base64.b64encode(inp).decode()
  # reverse the string except for potential = at the end
  out = reverse_equals(inp)
  return out

def insert_newlines(s):
  lines = []
  for i in range(0, len(s), 76):
    lines.append(s[i : i + 76])
  return '\n'.join(lines)

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument("-d", "--decode", required=False, action="store_true", help="Decode xmas code")
  parser.add_argument("-e", "--encode", required=False, action="store_true", help="Encode string to xmas code")
  parser.add_argument('input', nargs='?', help='optional input file')

  args = parser.parse_args()

  # assume decode when invoked as xmasdecode
  if os.path.basename(__file__) == "xmasdecode":
    args.decode = True
  else:
    if not args.decode:
      args.encode = True

  # decode
  if args.decode:
    if args.encode:
      print("Conflicting option -d and -e")
      sys.exit(1)

    # the input is ASCII text
    if args.input:
      outfile = args.input
      if outfile.endswith(SUFFIX):
        sys.stdout = open(outfile[:len(outfile) - len(SUFFIX)], 'wt')
      with open(args.input, 'rt') as f:
        inp = f.read()
    else:
      inp = sys.stdin.read()

    inp = strip_header(inp)
    sys.stdout.buffer.write(xmasdecode(inp))
  # encode
  else:
    # treat the inp as binary
    if args.input:
      outfile = args.input + SUFFIX
      sys.stdout = open(outfile, 'wt')
      with open(args.input, 'rb') as f:
        inp = f.read()
    else:
      inp = sys.stdin.buffer.read()

    # print the encoded file inp blocks
    out = insert_newlines(xmasencode(inp))
    out = prepend_header(out)
    print(out)

if __name__ == "__main__":
  main()
