#! /usr/bin/env python3
import sys
import pycdlib
from PySquashfsImage import SquashFsImage
from io import BytesIO

if len(sys.argv) != 3 and len(sys.argv) != 4:
    print('Usage: %s <path_to_print> <iso> OR <path_to_filesystem> <path_to_extract_in_filesystem> <iso>' % (sys.argv[0]))
    sys.exit(1)

iso = pycdlib.PyCdlib()

if len(sys.argv) == 3:
    iso.open(sys.argv[2])
    for child in iso.list_children(iso_path=sys.argv[1]):
        print(child.file_identifier().decode())
    iso.close()

if len(sys.argv) == 4:
    iso.open(sys.argv[3])
    print("Extract filesystem.squashfs from the iso ...")
    extracted = BytesIO()
    iso.get_file_from_iso_fp(extracted, iso_path=sys.argv[1])
    iso.close()
    file = open('filesystem.squashfs', 'wb')
    file.write(extracted.getvalue())
    file.close()

    print("Downloading %s ..." % sys.argv[2])
    image = SquashFsImage('filesystem.squashfs')
    for img in image.root.findAll():
        if img.getPath() == sys.argv[2]:
            with open('output', 'wb') as f:
                f.write(img.getContent())
    image.close()
