#-------------------------------------------------------------------------------
# Name:        import_quar.py
# purpose:     read a textfile with output from FortiOS CLI diag command,
#              create output command file for importing again
#              (e.g. to a different FGT)
#              usage: import_quar [filename] or
#                     import_quar < filename
#              output: quarlist.bcmd
# use the "import script" facility in Fortigate GUI to run the .bcmd
#
# created:     2023-02-10
# changed:     2026-01-22
#
# author:      Wolfgang Beneicke
# copyright:   (c) 2023 BENEICKE EDV-Beratung, Heidelberg, Germany
# Licence:     use as you please
#-------------------------------------------------------------------------------

import sys
import datetime as dt # for strptime(), datetime
import argparse


def get_args():
    ar = argparse.ArgumentParser(description='script Fortigate quarantine tool')
    ar.set_defaults(
        outfile = 'quarlist.bcmd',
        debug = False,
        duration = -1,
    )

    add = ar.add_argument    # an alias
    add('input', type=argparse.FileType('r'),
           help='read existing quarantine list from <%(dest)s>')
    add('-d','--debug', action='store_true', help='more verbose output')
    add('-o','--outfile', type=argparse.FileType('w'), help='output batch commands are in <%(dest)s>')
    add('-t','--duration', type=int, help='ban for <%(dest)s> seconds, 0 for permanent ban')

    args = ar.parse_args()
    return args


def main():
    args = get_args()

    # read the list from inputfile
    with args.input as f:
        qlist = f.readlines()

    with args.outfile as out:
        input_fmt = '%b %d %H:%M:%S %Y'
        read = written = 0
        permanent_ban = args.duration == 0  # global, for all IPs

        for line in qlist:
            read += 1
            if args.debug:
                print(f"{read}: {line}", end='')
            if not line[0].isdigit():        # skip header:
                if args.debug: print('    skipped')
                continue

            items = line.split()
            if not len(items) in (8, 12):
                continue
            # expected line content:
            # a.b.c.d    Sun Feb  5 18:32:44 2023 Sun Feb 12 18:32:44 2023 IPS
            # items:
            # 0: IPv4
            # 1- 5: start date= wd mon day time year
            # 6-10: expiry date=wd mon day time year
            # or "indefinite" (== "0" seconds) then len==8
            # 11: source (admin/dlp/ips/av/dos)

            # list output is exiry datestamp, input is duration in s from
            #   time of creation
            # so there is some calculation involved

            permanent_this = items[6].startswith('indef')

            IPv4 = items[0]
            by_whom = items[-1]
            by_whom = by_whom[:5].lower()  # max 5 chars for input, 'Administrative' in list
            # if source is misspelled when writing,
            # 'dlp' will be used instead by FortiOS

            if permanent_ban or permanent_this:
                duration = 0
            elif args.duration >= 0:
                duration = args.duration
            else:
                # leave out weekday
                expirydate = dt.datetime.strptime(' '.join(items[7:-1]), input_fmt)
                # calc duration from now to expiry date, in seconds
                delta = expirydate - dt.datetime.now()
                duration = int(delta.total_seconds())

            if duration >= 0:
                s = f"diag user quarantine add src4 {IPv4} {duration} {by_whom}"
                print(s, file=out)
                written += 1
                if args.debug:
                    print(f"{written}:    -> {s}")

    print(f'items written: {written}\nlines discarded: {read - written}')
    print('now import "{args.outfile.name}" into Fortigate as a batch command script')


if __name__ == '__main__':
    main()