utils/printing/paper at 91c2675f70c576972a92dda2f41286ed741f8abe · ocf/utils · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
#!/usr/bin/env python3
import argparse
import json
import sys
from datetime import datetime

from ocflib.account.search import user_exists
from ocflib.account.search import user_is_group
from ocflib.misc.shell import bold
from ocflib.misc.shell import green
from ocflib.misc.shell import red
from ocflib.misc.shell import yellow
from ocflib.misc.whoami import current_user
from ocflib.printing.quota import add_adjustment
from ocflib.printing.quota import get_connection
from ocflib.printing.quota import get_quota
from ocflib.printing.quota import Payload


def staff_creds():
    """Return dictionary of staff creds, if we can read them."""
    return json.load(open('/etc/ocfprinting.json'))


def view(args):
    with get_connection() as c:
        quota = get_quota(c, args.user)

    def highlight(num, text, warning_level):
        if num > warning_level:
            color = green
        elif 0 < num <= warning_level:
            color = yellow
        else:
            color = red
        return bold('  ' + color('→ {}{}'.format(num, text)))

    print('{} printing quota:'.format(bold(quota.user)))
    if not user_is_group(quota.user):
        print(highlight(quota.semesterly, ' remaining this semester', 30))
        print(highlight(quota.daily, ' remaining today', 5))
    else:
        print(bold(red("Group accounts can't print. Sorry!")))


def adjust(args, type):
    payload = Payload(
        user=args.user,
        time=datetime.now(),
        pages=args.pages,
        action=type,
        staffer=current_user(),
        reason=args.reason,
    )

    stringtype = 'Refund' if type == 'refund' else 'Forward'
    prompt = bold(stringtype + ' {} pages to {}? [yN] '.format(payload.pages, payload.user))
    if input(prompt) not in {'y', 'yes'}:
        print('Cancelled.')
        return
    try:
        credentials = staff_creds()
    except FileNotFoundError:
        print(red('Could not find the file for staff credentials.'))
        print(red('Are you running this on supernova?'))
        return 1
    with get_connection(**credentials) as c:
        add_adjustment(c, payload)
    print('Added.')


def main(argv=None):
    commands = {
        'view': view,
        'refund': adjust,
        'forward': adjust,
    }

    parser = argparse.ArgumentParser(description='View and manipulate page quotas.')
    subparsers = parser.add_subparsers(dest='command', help='command to run')
    subparsers.required = True

    parser_view = subparsers.add_parser('view', help="view a user's print quota")
    parser_view.add_argument('user', nargs='?', type=str, default=current_user())

    parser_refund = subparsers.add_parser('refund', help="refund a user's pages")
    parser_refund.add_argument('--pages', '-p', required=True, type=int)
    parser_refund.add_argument('--reason', '-r', required=True, type=str)
    parser_refund.add_argument('user', type=str)

    parser_forward = subparsers.add_parser('forward', help="allow a user to print above today's quota")
    parser_forward.add_argument('--pages', '-p', required=True, type=int)
    parser_forward.add_argument('--reason', '-r', required=True, type=str)
    parser_forward.add_argument('user', type=str)

    if len(argv or sys.argv[1:]) == 0:
        args = parser.parse_args(['view'])
    else:
        args = parser.parse_args(argv)

    if not user_exists(args.user):
        print(bold(red("The user {} doesn't exist.".format(args.user))))
        return 1

    if args.command == 'refund' or args.command == 'forward':
        return commands[args.command](args, args.command)
    else:
        return commands[args.command](args)


if __name__ == '__main__':
    sys.exit(main())