<s> import os <EOL> import re <EOL> import json <EOL> import string <EOL> import datetime <EOL> from base64 import b32decode <EOL> from six import StringIO <EOL> import unittest <EOL> from django . test import TestCase , Client <EOL> from django . core . urlresolvers import reverse <EOL> from django . contrib . auth import get_user_model , SESSION_KEY <EOL> from django . utils . crypto import get_random_string <EOL> from django . utils import timezone <EOL> from django . core . management import call_command <EOL> from django_u2f . forms import BackupCodeForm , TOTPForm <EOL> from django_u2f import oath <EOL> from django_u2f . models import TOTPDevice <EOL> User = get_user_model ( ) <EOL> class TwoFactorTest ( TestCase ) : <EOL> def setUp ( self ) : <EOL> self . client = Client ( ) <EOL> self . user = User . objects . create_superuser ( <EOL> username = '<STR_LIT:test>' , <EOL> email = '<STR_LIT>' , <EOL> password = '<STR_LIT>' , <EOL> ) <EOL> self . login_url = reverse ( '<STR_LIT>' ) <EOL> def login ( self ) : <EOL> return self . client . post ( self . login_url , { <EOL> '<STR_LIT:username>' : '<STR_LIT:test>' , <EOL> '<STR_LIT:password>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } ) <EOL> def enable_backupcode ( self ) : <EOL> code = get_random_string ( length = <NUM_LIT:6> , allowed_chars = string . digits ) <EOL> self . user . backup_codes . create ( code = code ) <EOL> return code <EOL> class U2FTest ( TwoFactorTest ) : <EOL> def enable_u2f ( self ) : <EOL> self . user . u2f_keys . create ( <EOL> public_key = '<STR_LIT>' , <EOL> key_handle = '<STR_LIT>' , <EOL> app_id = '<STR_LIT>' , <EOL> ) <EOL> def set_challenge ( self ) : <EOL> session = self . client . session <EOL> session [ '<STR_LIT>' ] = [ <EOL> { <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT:version>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> } <EOL> ] <EOL> session . save ( ) <EOL> return { <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> } <EOL> def set_add_key ( self ) : <EOL> session = self . client . session <EOL> session [ '<STR_LIT>' ] = { <EOL> '<STR_LIT>' : "<STR_LIT>" , <EOL> '<STR_LIT:version>' : "<STR_LIT>" , <EOL> '<STR_LIT>' : "<STR_LIT>" , <EOL> } <EOL> session . save ( ) <EOL> return { <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT:version>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> } <EOL> class TestU2F ( U2FTest ) : <EOL> def test_normal_login ( self ) : <EOL> r = self . login ( ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( '<STR_LIT>' ) ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> def test_u2f_login ( self ) : <EOL> self . enable_u2f ( ) <EOL> r = self . login ( ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertIn ( reverse ( '<STR_LIT>' ) , r [ '<STR_LIT:location>' ] ) <EOL> device_response = self . set_challenge ( ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT>' : json . dumps ( device_response ) , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertEquals ( r . status_code , <NUM_LIT> ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( '<STR_LIT>' ) ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> def test_failed_u2f_login ( self ) : <EOL> self . enable_u2f ( ) <EOL> r = self . login ( ) <EOL> device_response = self . set_challenge ( ) <EOL> device_response [ '<STR_LIT>' ] = '<STR_LIT:a>' + device_response [ '<STR_LIT>' ] [ <NUM_LIT:1> : ] <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT>' : json . dumps ( device_response ) , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertContains ( r , '<STR_LIT>' ) <EOL> def test_verify_when_not_logged_in ( self ) : <EOL> r = self . client . get ( reverse ( '<STR_LIT>' ) ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( self . login_url ) ) <EOL> def test_add_key ( self ) : <EOL> self . login ( ) <EOL> url = reverse ( '<STR_LIT>' ) <EOL> r = self . client . get ( url ) <EOL> self . assertContains ( r , '<STR_LIT>' ) <EOL> device_response = self . set_add_key ( ) <EOL> r = self . client . post ( url , { <EOL> '<STR_LIT>' : json . dumps ( device_response ) , <EOL> } ) <EOL> added_key_obj = self . user . u2f_keys . get ( ) <EOL> self . assertEqual ( added_key_obj . public_key , '<STR_LIT>' ) <EOL> self . assertEqual ( added_key_obj . key_handle , '<STR_LIT>' ) <EOL> self . assertEqual ( added_key_obj . app_id , '<STR_LIT>' ) <EOL> def test_key_delete ( self ) : <EOL> other_user = User . objects . create_superuser ( <EOL> username = '<STR_LIT:abc>' , <EOL> email = '<STR_LIT>' , <EOL> password = '<STR_LIT>' , <EOL> ) <EOL> other_user . u2f_keys . create ( ) <EOL> self . enable_u2f ( ) <EOL> self . client . login ( username = '<STR_LIT:test>' , password = '<STR_LIT>' ) <EOL> r = self . client . post ( reverse ( '<STR_LIT>' ) , { <EOL> '<STR_LIT>' : self . user . u2f_keys . get ( ) . pk , <EOL> '<STR_LIT>' : '<STR_LIT:1>' , <EOL> } ) <EOL> self . assertFalse ( self . user . u2f_keys . exists ( ) ) <EOL> r = self . client . post ( reverse ( '<STR_LIT>' ) , { <EOL> '<STR_LIT>' : other_user . u2f_keys . get ( ) . pk , <EOL> '<STR_LIT>' : '<STR_LIT:1>' , <EOL> } ) <EOL> self . assertTrue ( other_user . u2f_keys . exists ( ) ) <EOL> self . assertEqual ( r . status_code , <NUM_LIT> ) <EOL> class TestAdminLogin ( TwoFactorTest ) : <EOL> def setUp ( self ) : <EOL> super ( TestAdminLogin , self ) . setUp ( ) <EOL> self . admin_url = reverse ( '<STR_LIT>' ) <EOL> r = self . client . get ( self . admin_url ) <EOL> if r . status_code == <NUM_LIT> : <EOL> self . login_url = r [ '<STR_LIT:location>' ] <EOL> else : <EOL> self . login_url = self . admin_url <EOL> def test_admin_template_rendered ( self ) : <EOL> r = self . client . get ( self . login_url ) <EOL> self . assertEqual ( r . templates [ <NUM_LIT:0> ] . name , '<STR_LIT>' ) <EOL> def test_login_with_u2f ( self ) : <EOL> code = self . enable_backupcode ( ) <EOL> r = self . client . post ( self . login_url , { <EOL> '<STR_LIT:username>' : '<STR_LIT:test>' , <EOL> '<STR_LIT:password>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertIn ( reverse ( '<STR_LIT>' ) , r [ '<STR_LIT:location>' ] ) <EOL> verify_key_response = self . client . get ( r [ '<STR_LIT:location>' ] ) <EOL> self . assertContains ( verify_key_response , '<STR_LIT>' ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT:code>' : code , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( self . admin_url ) ) <EOL> def test_login_without_u2f ( self ) : <EOL> r = self . client . post ( self . login_url , { <EOL> '<STR_LIT:username>' : '<STR_LIT:test>' , <EOL> '<STR_LIT:password>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : <NUM_LIT:1> , <EOL> '<STR_LIT>' : self . admin_url , <EOL> } ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( self . admin_url ) ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> class TestBackupCode ( TwoFactorTest ) : <EOL> def test_validation_error ( self ) : <EOL> self . enable_backupcode ( ) <EOL> r = self . login ( ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertIn ( reverse ( '<STR_LIT>' ) , r [ '<STR_LIT:location>' ] ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertContains ( r , '<STR_LIT>' ) <EOL> def test_code_required_and_login ( self ) : <EOL> code = self . enable_backupcode ( ) <EOL> r = self . login ( ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertIn ( reverse ( '<STR_LIT>' ) , r [ '<STR_LIT:location>' ] ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT:code>' : code , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( '<STR_LIT>' ) ) <EOL> self . assertFalse ( self . user . backup_codes . filter ( code = code ) . exists ( ) ) <EOL> def test_incorrect_code ( self ) : <EOL> self . enable_backupcode ( ) <EOL> r = self . login ( ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT:code>' : '<STR_LIT:abc>' , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertContains ( r , BackupCodeForm . INVALID_ERROR_MESSAGE ) <EOL> def test_add_backup_codes ( self ) : <EOL> self . login ( ) <EOL> r = self . client . post ( reverse ( '<STR_LIT>' ) ) <EOL> self . assertEqual ( r . status_code , <NUM_LIT> ) <EOL> self . assertEqual ( self . user . backup_codes . count ( ) , <NUM_LIT:10> ) <EOL> def test_list_backup_codes ( self ) : <EOL> self . login ( ) <EOL> self . user . backup_codes . create ( code = '<STR_LIT>' ) <EOL> otheruser = User . objects . create_superuser ( <EOL> username = '<STR_LIT>' , <EOL> email = '<STR_LIT>' , <EOL> password = '<STR_LIT>' , <EOL> ) <EOL> otheruser . backup_codes . create ( code = '<STR_LIT>' ) <EOL> r = self . client . get ( reverse ( '<STR_LIT>' ) ) <EOL> self . assertContains ( r , '<STR_LIT>' ) <EOL> self . assertNotContains ( r , '<STR_LIT>' ) <EOL> def test_addbackupcode ( self ) : <EOL> out = StringIO ( ) <EOL> call_command ( '<STR_LIT>' , self . user . get_username ( ) , stdout = out ) <EOL> code = out . getvalue ( ) . strip ( ) <EOL> self . assertTrue ( self . user . backup_codes . filter ( code = code ) . exists ( ) ) <EOL> call_command ( '<STR_LIT>' , self . user . get_username ( ) , code = '<STR_LIT:foo>' , stdout = out ) <EOL> self . assertTrue ( self . user . backup_codes . filter ( code = '<STR_LIT:foo>' ) . exists ( ) ) <EOL> def test_login_while_already_logged_in ( self ) : <EOL> User . objects . create_superuser ( <EOL> username = '<STR_LIT>' , <EOL> email = '<STR_LIT>' , <EOL> password = '<STR_LIT>' , <EOL> ) <EOL> r = self . client . post ( self . login_url , { <EOL> '<STR_LIT:username>' : '<STR_LIT>' , <EOL> '<STR_LIT:password>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' <EOL> } ) <EOL> assert r . status_code == <NUM_LIT> <EOL> code = self . enable_backupcode ( ) <EOL> r = self . login ( ) <EOL> self . assertIn ( reverse ( '<STR_LIT>' ) , r [ '<STR_LIT:location>' ] ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT:code>' : code , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( '<STR_LIT>' ) ) <EOL> self . assertFalse ( self . user . backup_codes . filter ( code = code ) . exists ( ) ) <EOL> class TestTOTP ( U2FTest ) : <EOL> def enable_totp ( self ) : <EOL> key = os . urandom ( <NUM_LIT:20> ) <EOL> self . user . totp_devices . create ( <EOL> key = key , <EOL> ) <EOL> return key <EOL> def test_login ( self ) : <EOL> key = self . enable_totp ( ) <EOL> r = self . login ( ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertIn ( reverse ( '<STR_LIT>' ) , r [ '<STR_LIT:location>' ] ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT>' : oath . totp ( key , timezone . now ( ) ) , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> self . assertTrue ( r [ '<STR_LIT:location>' ] . endswith ( '<STR_LIT>' ) ) <EOL> def test_token_cant_be_used_twice ( self ) : <EOL> key = self . enable_totp ( ) <EOL> r = self . login ( ) <EOL> token = oath . totp ( key , timezone . now ( ) ) , <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT>' : token , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertEqual ( str ( self . client . session [ SESSION_KEY ] ) , str ( self . user . id ) ) <EOL> self . client . logout ( ) <EOL> r = self . login ( ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT>' : token , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertContains ( r , TOTPForm . INVALID_ERROR_MESSAGE ) <EOL> def test_incorrect_code ( self ) : <EOL> key = self . enable_totp ( ) <EOL> r = self . login ( ) <EOL> r = self . client . post ( r [ '<STR_LIT:location>' ] , { <EOL> '<STR_LIT>' : oath . totp ( key , timezone . now ( ) + datetime . timedelta ( seconds = <NUM_LIT> ) ) , <EOL> '<STR_LIT:type>' : '<STR_LIT>' , <EOL> } ) <EOL> self . assertNotIn ( SESSION_KEY , self . client . session ) <EOL> self . assertContains ( r , TOTPForm . INVALID_ERROR_MESSAGE ) <EOL> def _extract_key ( self , response ) : <EOL> return re . search ( '<STR_LIT>' , response . content . decode ( '<STR_LIT:utf-8>' ) ) . group ( <NUM_LIT:1> ) <EOL> def test_add_device ( self ) : <EOL> self . login ( ) <EOL> url = reverse ( '<STR_LIT>' ) <EOL> r = self . client . get ( url ) <EOL> self . assertContains ( r , '<STR_LIT>' ) <EOL> base32_key = self . _extract_key ( r ) <EOL> key = b32decode ( base32_key ) <EOL> r = self . client . post ( url , { <EOL> '<STR_LIT>' : base32_key , <EOL> '<STR_LIT>' : oath . totp ( key , timezone . now ( ) ) , <EOL> } ) <EOL> self . assertEqual ( r . status_code , <NUM_LIT> ) <EOL> self . assertTrue ( self . user . totp_devices . filter ( key = key ) . exists ( ) ) <EOL> def test_add_device_incorrect_token ( self ) : <EOL> self . login ( ) <EOL> url = reverse ( '<STR_LIT>' ) <EOL> r = self . client . get ( url ) <EOL> base32_key = self . _extract_key ( r ) <EOL> key = b32decode ( base32_key ) <EOL> r = self . client . post ( url , { <EOL> '<STR_LIT>' : base32_key , <EOL> '<STR_LIT>' : oath . totp ( key , timezone . now ( ) + datetime . timedelta ( seconds = <NUM_LIT> ) ) , <EOL> } ) <EOL> self . assertEqual ( r . status_code , <NUM_LIT:200> ) <EOL> self . assertContains ( r , TOTPForm . INVALID_ERROR_MESSAGE ) <EOL> def test_delete_device ( self ) : <EOL> self . login ( ) <EOL> self . enable_totp ( ) <EOL> r = self . client . post ( reverse ( '<STR_LIT>' ) , { <EOL> '<STR_LIT>' : self . user . totp_devices . get ( ) . pk , <EOL> '<STR_LIT>' : '<STR_LIT:1>' , <EOL> } ) <EOL> self . assertEqual ( r . status_code , <NUM_LIT> ) <EOL> self . assertFalse ( self . user . totp_devices . exists ( ) ) <EOL> def test_slop ( self ) : <EOL> key = os . urandom ( <NUM_LIT:20> ) <EOL> device = TOTPDevice ( key = key ) <EOL> now = timezone . now ( ) <EOL> self . assertTrue ( device . validate_token ( oath . totp ( key , now - datetime . timedelta ( seconds = <NUM_LIT:30> ) ) ) ) <EOL> self . assertTrue ( device . validate_token ( oath . totp ( key , now ) ) ) <EOL> self . assertTrue ( device . validate_token ( oath . totp ( key , now + datetime . timedelta ( seconds = <NUM_LIT:30> ) ) ) ) </s>
