diff --git a/patch-vbmeta.py b/patch-vbmeta.py index 0921e45..0b7f4e3 100644 --- a/patch-vbmeta.py +++ b/patch-vbmeta.py @@ -10,35 +10,35 @@ FLAGS_TO_SET = b'\x03' -### main program - -# if an argument is not provided -if len(sys.argv) != 2: - sys.exit("Usage: python ./" + os.path.basename(__file__) + " ") - -# try reading the file with read/write to make sure it exists -FILE = sys.argv[1] - -try: - fd = os.open(FILE, os.O_RDWR) -except OSError: - sys.exit("Error reading file: " + FILE + ". File not modified. Exiting...") - -# making sure it's a vbmeta image by reading the magic bytes at the start of the file -magic = os.read(fd, AVB_MAGIC_LEN) - -if (magic != AVB_MAGIC): - fd.close() - sys.exit("Error: The provided image is not a valid vbmeta image. File not modified. Exiting...") - -# set the disable-verity and disable-verification flags at offset 123 -try: - os.lseek(fd, FLAGS_OFFSET, os.SEEK_SET) - os.write(fd, FLAGS_TO_SET) -except OSError: - fd.close() - sys.exit("Error: Failed when patching the vbmeta image. Exiting...") - -# end of program -os.close(fd) -print("Patching successful.") \ No newline at end of file +if __name__ == "__main__": + + # if an argument is not provided + if len(sys.argv) != 2: + sys.exit("Usage: python ./" + os.path.basename(__file__) + " ") + + # try reading the file with read/write to make sure it exists + FILE = sys.argv[1] + + try: + fd = os.open(FILE, os.O_RDWR) + except OSError: + sys.exit("Error reading file: " + FILE + ". File not modified. Exiting...") + + # making sure it's a vbmeta image by reading the magic bytes at the start of the file + magic = os.read(fd, AVB_MAGIC_LEN) + + if (magic != AVB_MAGIC): + fd.close() + sys.exit("Error: The provided image is not a valid vbmeta image. File not modified. Exiting...") + + # set the disable-verity and disable-verification flags at offset 123 + try: + os.lseek(fd, FLAGS_OFFSET, os.SEEK_SET) + os.write(fd, FLAGS_TO_SET) + except OSError: + fd.close() + sys.exit("Error: Failed when patching the vbmeta image. Exiting...") + + # end of program + os.close(fd) + print("Patching successful.") \ No newline at end of file